Merge branch 'main' into feat/linuxArm

Author: kratman

CMakeLists.txt

@@ -57,27 +57,6 @@ if (${PYBAMM_IDAKLU_EXPR_CASADI} STREQUAL "ON")
     )
 endif ()
 
-# Check IREE build flag
-if (NOT DEFINED PYBAMM_IDAKLU_EXPR_IREE)
-    set(PYBAMM_IDAKLU_EXPR_IREE OFF)
-endif ()
-message("PYBAMM_IDAKLU_EXPR_IREE: ${PYBAMM_IDAKLU_EXPR_IREE}")
-
-# IREE (MLIR expression evaluation) PyBaMM source files
-set(IDAKLU_EXPR_IREE_SOURCE_FILES "")
-if (${PYBAMM_IDAKLU_EXPR_IREE} STREQUAL "ON")
-    add_compile_definitions(IREE_ENABLE)
-    # Source file list
-    set(IDAKLU_EXPR_IREE_SOURCE_FILES
-            src/pybammsolvers/idaklu_source/Expressions/IREE/iree_jit.cpp
-            src/pybammsolvers/idaklu_source/Expressions/IREE/iree_jit.hpp
-            src/pybammsolvers/idaklu_source/Expressions/IREE/IREEFunctions.cpp
-            src/pybammsolvers/idaklu_source/Expressions/IREE/IREEFunctions.hpp
-            src/pybammsolvers/idaklu_source/Expressions/IREE/ModuleParser.cpp
-            src/pybammsolvers/idaklu_source/Expressions/IREE/ModuleParser.hpp
-    )
-endif ()
-
 # The complete (all dependencies) sources list should be mirrored in setup.py
 pybind11_add_module(idaklu
         # pybind11 interface
@@ -113,7 +92,6 @@ pybind11_add_module(idaklu
         src/pybammsolvers/idaklu_source/Expressions/Base/ExpressionTypes.hpp
         # IDAKLU expressions - concrete implementations
         ${IDAKLU_EXPR_CASADI_SOURCE_FILES}
-        ${IDAKLU_EXPR_IREE_SOURCE_FILES}
 )
 
 if (NOT DEFINED USE_PYTHON_CASADI)
@@ -184,16 +162,3 @@ else ()
 endif ()
 include_directories(${SuiteSparse_INCLUDE_DIRS})
 target_link_libraries(idaklu PRIVATE ${SuiteSparse_LIBRARIES})
-
-# IREE (MLIR compiler and runtime library) build settings
-if (${PYBAMM_IDAKLU_EXPR_IREE} STREQUAL "ON")
-    set(IREE_BUILD_COMPILER ON)
-    set(IREE_BUILD_TESTS OFF)
-    set(IREE_BUILD_SAMPLES OFF)
-    add_subdirectory(iree EXCLUDE_FROM_ALL)
-    set(IREE_COMPILER_ROOT "${CMAKE_CURRENT_SOURCE_DIR}/iree/compiler")
-    target_include_directories(idaklu SYSTEM PRIVATE "${IREE_COMPILER_ROOT}/bindings/c/iree/compiler")
-    target_compile_options(idaklu PRIVATE ${IREE_DEFAULT_COPTS})
-    target_link_libraries(idaklu PRIVATE iree_compiler_bindings_c_loader)
-    target_link_libraries(idaklu PRIVATE iree_runtime_runtime)
-endif ()

noxfile.py

@@ -1,8 +1,6 @@
 import nox
 import os
 import sys
-import warnings
-import platform
 from pathlib import Path
 
 
@@ -13,55 +11,13 @@
 else:
     nox.options.sessions = ["unit"]
 
-
-def set_iree_state():
-    """
-    Check if IREE is enabled and set the environment variable accordingly.
-
-    Returns
-    -------
-    str
-        "ON" if IREE is enabled, "OFF" otherwise.
-
-    """
-    state = "ON" if os.getenv("PYBAMM_IDAKLU_EXPR_IREE", "OFF") == "ON" else "OFF"
-    if state == "ON":
-        if sys.platform == "win32":
-            warnings.warn(
-                (
-                    "IREE is not enabled on Windows yet. "
-                    "Setting PYBAMM_IDAKLU_EXPR_IREE=OFF."
-                ),
-                stacklevel=2,
-            )
-            return "OFF"
-        if sys.platform == "darwin":
-            # iree-compiler is currently only available as a wheel on macOS 13 (or
-            # higher) and Python version 3.11
-            mac_ver = int(platform.mac_ver()[0].split(".")[0])
-            if (not sys.version_info[:2] == (3, 11)) or mac_ver < 13:
-                warnings.warn(
-                    (
-                        "IREE is only supported on MacOS 13 (or higher) and Python"
-                        "version 3.11. Setting PYBAMM_IDAKLU_EXPR_IREE=OFF."
-                    ),
-                    stacklevel=2,
-                )
-                return "OFF"
-    return state
-
-
 homedir = Path(__file__)
 PYBAMM_ENV = {
     "LD_LIBRARY_PATH": f"{homedir}/.idaklu/lib",
     "PYTHONIOENCODING": "utf-8",
     "MPLBACKEND": "Agg",
     # Expression evaluators (...EXPR_CASADI cannot be fully disabled at this time)
     "PYBAMM_IDAKLU_EXPR_CASADI": os.getenv("PYBAMM_IDAKLU_EXPR_CASADI", "ON"),
-    "PYBAMM_IDAKLU_EXPR_IREE": set_iree_state(),
-    "IREE_INDEX_URL": os.getenv(
-        "IREE_INDEX_URL", "https://iree.dev/pip-release-links.html"
-    ),
 }
 VENV_DIR = Path("./venv").resolve()
 
@@ -89,29 +45,6 @@ def run_pybamm_requires(session):
     set_environment_variables(PYBAMM_ENV, session=session)
     if sys.platform != "win32":
         session.run("python", "install_KLU_Sundials.py", *session.posargs)
-        if PYBAMM_ENV.get("PYBAMM_IDAKLU_EXPR_IREE") == "ON" and not os.path.exists(
-            "./iree"
-        ):
-            session.run(
-                "git",
-                "clone",
-                "--depth=1",
-                "--recurse-submodules",
-                "--shallow-submodules",
-                "--branch=candidate-20240507.886",
-                "https://github.com/openxla/iree",
-                "iree/",
-                external=True,
-            )
-            with session.chdir("iree"):
-                session.run(
-                    "git",
-                    "submodule",
-                    "update",
-                    "--init",
-                    "--recursive",
-                    external=True,
-                )
     else:
         session.error("nox -s idaklu-requires is only available on Linux & macOS.")
 
@@ -122,12 +55,4 @@ def run_unit(session):
     session.install("setuptools", silent=False)
     session.install("casadi==3.6.7", silent=False)
     session.install(".[dev]", silent=False)
-    if PYBAMM_ENV.get("PYBAMM_IDAKLU_EXPR_IREE") == "ON":
-        # See comments in 'dev' session
-        session.install(
-            ".[iree]",
-            "--find-links",
-            PYBAMM_ENV.get("IREE_INDEX_URL"),
-            silent=False,
-        )
     session.run("pytest", "tests")

setup.py

@@ -92,13 +92,11 @@ def run(self):
 
         build_type = os.getenv("PYBAMM_CPP_BUILD_TYPE", "Release")
         idaklu_expr_casadi = os.getenv("PYBAMM_IDAKLU_EXPR_CASADI", "ON")
-        idaklu_expr_iree = os.getenv("PYBAMM_IDAKLU_EXPR_IREE", "OFF")
         cmake_args = [
             f"-DCMAKE_BUILD_TYPE={build_type}",
             f"-DPYTHON_EXECUTABLE={sys.executable}",
             "-DUSE_PYTHON_CASADI={}".format("TRUE" if use_python_casadi else "FALSE"),
             f"-DPYBAMM_IDAKLU_EXPR_CASADI={idaklu_expr_casadi}",
-            f"-DPYBAMM_IDAKLU_EXPR_IREE={idaklu_expr_iree}",
         ]
         if self.suitesparse_root:
             cmake_args.append(
@@ -244,14 +242,6 @@ def run(self):
             "src/pybammsolvers/idaklu_source/Expressions/Base/ExpressionSparsity.hpp",
             "src/pybammsolvers/idaklu_source/Expressions/Casadi/CasadiFunctions.cpp",
             "src/pybammsolvers/idaklu_source/Expressions/Casadi/CasadiFunctions.hpp",
-            "src/pybammsolvers/idaklu_source/Expressions/IREE/IREEBaseFunction.hpp",
-            "src/pybammsolvers/idaklu_source/Expressions/IREE/IREEFunction.hpp",
-            "src/pybammsolvers/idaklu_source/Expressions/IREE/IREEFunctions.cpp",
-            "src/pybammsolvers/idaklu_source/Expressions/IREE/IREEFunctions.hpp",
-            "src/pybammsolvers/idaklu_source/Expressions/IREE/iree_jit.cpp",
-            "src/pybammsolvers/idaklu_source/Expressions/IREE/iree_jit.hpp",
-            "src/pybammsolvers/idaklu_source/Expressions/IREE/ModuleParser.cpp",
-            "src/pybammsolvers/idaklu_source/Expressions/IREE/ModuleParser.hpp",
             "src/pybammsolvers/idaklu_source/idaklu_solver.hpp",
             "src/pybammsolvers/idaklu_source/IDAKLUSolver.cpp",
             "src/pybammsolvers/idaklu_source/IDAKLUSolver.hpp",

src/pybammsolvers/idaklu.cpp

@@ -15,10 +15,6 @@
 #include "idaklu_source/common.hpp"
 #include "idaklu_source/Expressions/Casadi/CasadiFunctions.hpp"
 
-#ifdef IREE_ENABLE
-#include "idaklu_source/Expressions/IREE/IREEFunctions.hpp"
-#endif
-
 
 casadi::Function generate_casadi_function(const std::string &data)
 {
@@ -96,34 +92,6 @@ PYBIND11_MODULE(idaklu, m)
     py::arg("shape"),
     py::return_value_policy::take_ownership);
 
-#ifdef IREE_ENABLE
-  m.def("create_iree_solver_group", &create_idaklu_solver_group<IREEFunctions>,
-    "Create a group of iree idaklu solver objects",
-    py::arg("number_of_states"),
-    py::arg("number_of_parameters"),
-    py::arg("rhs_alg"),
-    py::arg("jac_times_cjmass"),
-    py::arg("jac_times_cjmass_colptrs"),
-    py::arg("jac_times_cjmass_rowvals"),
-    py::arg("jac_times_cjmass_nnz"),
-    py::arg("jac_bandwidth_lower"),
-    py::arg("jac_bandwidth_upper"),
-    py::arg("jac_action"),
-    py::arg("mass_action"),
-    py::arg("sens"),
-    py::arg("events"),
-    py::arg("number_of_events"),
-    py::arg("rhs_alg_id"),
-    py::arg("atol"),
-    py::arg("rtol"),
-    py::arg("inputs"),
-    py::arg("var_fcns"),
-    py::arg("dvar_dy_fcns"),
-    py::arg("dvar_dp_fcns"),
-    py::arg("options"),
-    py::return_value_policy::take_ownership);
-#endif
-
   m.def("generate_function", &generate_casadi_function,
     "Generate a casadi function",
     py::arg("string"),
@@ -174,20 +142,6 @@ PYBIND11_MODULE(idaklu, m)
 
   py::class_<casadi::Function>(m, "Function");
 
-#ifdef IREE_ENABLE
-  py::class_<IREEBaseFunctionType>(m, "IREEBaseFunctionType")
-    .def(py::init<>())
-    .def_readwrite("mlir", &IREEBaseFunctionType::mlir)
-    .def_readwrite("kept_var_idx", &IREEBaseFunctionType::kept_var_idx)
-    .def_readwrite("nnz", &IREEBaseFunctionType::nnz)
-    .def_readwrite("numel", &IREEBaseFunctionType::numel)
-    .def_readwrite("col", &IREEBaseFunctionType::col)
-    .def_readwrite("row", &IREEBaseFunctionType::row)
-    .def_readwrite("pytree_shape", &IREEBaseFunctionType::pytree_shape)
-    .def_readwrite("pytree_sizes", &IREEBaseFunctionType::pytree_sizes)
-    .def_readwrite("n_args", &IREEBaseFunctionType::n_args);
-#endif
-
   py::class_<Solution>(m, "solution")
     .def_readwrite("t", &Solution::t)
     .def_readwrite("y", &Solution::y)