faster-diffbloch 0.1.0__py3-none-any.whl

faster_diffbloch/__init__.py

@@ -0,0 +1,14 @@
+"""faster-diffBloch: Metal GPU and CPU acceleration for diffBloch."""
+
+from __future__ import annotations
+
+from .backend import enable, disable, matrix_exp, matrix_exp_backward, faster_propagate
+
+__version__ = "0.1.0"
+__all__ = [
+    "enable",
+    "disable",
+    "matrix_exp",
+    "matrix_exp_backward",
+    "faster_propagate",
+]

faster_diffbloch/backend.py

@@ -0,0 +1,149 @@
+"""Flow and Metal acceleration backend for diffBloch."""
+
+from __future__ import annotations
+
+import ctypes
+from typing import Literal
+import numpy as np
+import torch
+
+from .builder import build_and_load_library
+
+_LIB: ctypes.CDLL | None = None
+_CURRENT_DEVICE: Literal["cpu", "gpu"] = "gpu"
+
+
+def get_library() -> ctypes.CDLL:
+    global _LIB
+    if _LIB is None:
+        _LIB = build_and_load_library()
+    return _LIB
+
+
+def matrix_exp(a: np.ndarray, device: str = "gpu") -> np.ndarray:
+    lib = get_library()
+    a_c = np.ascontiguousarray(a, dtype=np.complex64)
+    out = np.empty_like(a_c)
+    n = a_c.shape[-1]
+    batch = int(a_c.size // (n * n))
+    ptr_in = a_c.ctypes.data_as(ctypes.c_void_p)
+    ptr_out = out.ctypes.data_as(ctypes.c_void_p)
+
+    if device == "gpu":
+        fn = getattr(lib, "bridge_matrix_exp_gpu_ptr_c64_i32_i32_ptr_c64", None)
+        if fn:
+            fn.argtypes = [ctypes.c_void_p, ctypes.c_int32, ctypes.c_int32, ctypes.c_void_p]
+            fn.restype = None
+            fn(ptr_in, batch, n, ptr_out)
+            return out
+
+    fn = getattr(lib, "bridge_matrix_exp_ptr_c64_i32_i32_ptr_c64")
+    fn.argtypes = [ctypes.c_void_p, ctypes.c_int32, ctypes.c_int32, ctypes.c_void_p]
+    fn.restype = None
+    fn(ptr_in, batch, n, ptr_out)
+    return out
+
+
+def matrix_exp_backward(a: np.ndarray, ebar: np.ndarray, dense: bool = False, device: str = "gpu") -> np.ndarray:
+    lib = get_library()
+    a_c = np.ascontiguousarray(a, dtype=np.complex64)
+    ebar_c = np.ascontiguousarray(ebar, dtype=np.complex64)
+    out = np.empty_like(a_c)
+    n = a_c.shape[-1]
+    batch = int(a_c.size // (n * n))
+    dense_val = 1 if dense else 0
+
+    ptr_a = a_c.ctypes.data_as(ctypes.c_void_p)
+    ptr_ebar = ebar_c.ctypes.data_as(ctypes.c_void_p)
+    ptr_out = out.ctypes.data_as(ctypes.c_void_p)
+
+    if device == "gpu":
+        fn = getattr(lib, "bridge_matrix_exp_backward_gpu_ptr_c64_ptr_c64_i32_i32_i32_ptr_c64", None)
+        if fn:
+            fn.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_int32, ctypes.c_int32, ctypes.c_int32, ctypes.c_void_p]
+            fn.restype = None
+            fn(ptr_a, ptr_ebar, dense_val, batch, n, ptr_out)
+            return out
+
+    fn = getattr(lib, "bridge_matrix_exp_backward_ptr_c64_ptr_c64_i32_i32_i32_ptr_c64")
+    fn.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_int32, ctypes.c_int32, ctypes.c_int32, ctypes.c_void_p]
+    fn.restype = None
+    fn(ptr_a, ptr_ebar, dense_val, batch, n, ptr_out)
+    return out
+
+
+class FasterMatrixExp(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, m: torch.Tensor) -> torch.Tensor:
+        source = m.detach().resolve_conj().contiguous()
+        exp_np = matrix_exp(source.cpu().numpy(), device=_CURRENT_DEVICE)
+        ctx.save_for_backward(source)
+        return torch.from_numpy(exp_np).to(device=m.device)
+
+    @staticmethod
+    @torch.autograd.function.once_differentiable
+    def backward(ctx, grad_output: torch.Tensor) -> torch.Tensor:
+        (source,) = ctx.saved_tensors
+        cotangent = grad_output.resolve_conj().contiguous().to(torch.complex64)
+        pullback_np = matrix_exp_backward(source.cpu().numpy(), cotangent.cpu().numpy(), dense=False, device=_CURRENT_DEVICE)
+        return torch.from_numpy(pullback_np).to(device=grad_output.device)
+
+
+def faster_propagate_matrix_exp(system, thicknesses, *, max_batch=None):
+    from diffBloch.core.solver import _complex_operator
+    a = _complex_operator(system.a).to(torch.complex64)
+    psi0 = system.psi0.to(dtype=a.dtype, device=a.device)
+    scalars = (1j * torch.pi * thicknesses / system.k_n).to(a.dtype)
+    if max_batch is None:
+        transfer = FasterMatrixExp.apply(a.unsqueeze(-3) * scalars[:, None, None])
+        return (transfer @ psi0.unsqueeze(-1)).squeeze(-1)
+    n = a.shape[-1]
+    a_flat = a.reshape(-1, n, n)
+    n_batch, n_thick = a_flat.shape[0], scalars.shape[0]
+    total = n_batch * n_thick
+    amplitudes = []
+    for start in range(0, total, max_batch):
+        flat = torch.arange(start, min(total, start + max_batch), device=a.device)
+        block = a_flat[flat // n_thick] * scalars[flat % n_thick][:, None, None]
+        amplitudes.append((FasterMatrixExp.apply(block) @ psi0.unsqueeze(-1)).squeeze(-1))
+    return torch.cat(amplitudes, dim=0).reshape(*a.shape[:-2], n_thick, n)
+
+
+def faster_propagate(system, thicknesses, *, method="matrix_exp", max_batch=None):
+    from diffBloch.core.solver import _propagate_bloch_eigen
+    if max_batch is not None and max_batch < 1:
+        raise ValueError(f"max_batch must be a positive integer or None, got {max_batch}")
+    t = torch.as_tensor(thicknesses, dtype=torch.float32, device=system.a.device)
+    if t.ndim == 0:
+        t = t.reshape(1)
+    if t.ndim != 1:
+        raise ValueError("thicknesses must be a scalar or 1-D sequence")
+    if method == "matrix_exp":
+        return faster_propagate_matrix_exp(system, t, max_batch=max_batch)
+    if method == "bloch_eigen":
+        return _propagate_bloch_eigen(system, t)
+    raise ValueError(f"method must be 'matrix_exp' or 'bloch_eigen', got {method!r}")
+
+
+def enable(device: Literal["cpu", "gpu"] = "gpu") -> None:
+    """Inject faster-diffbloch acceleration into diffBloch runtime."""
+    global _CURRENT_DEVICE
+    _CURRENT_DEVICE = device
+    import importlib
+    try:
+        solver = importlib.import_module("diffBloch.core.solver")
+        solver._propagate_matrix_exp = faster_propagate_matrix_exp
+        solver.propagate = faster_propagate
+        for mod_name in ("diffBloch.core", "diffBloch.engine.forward"):
+            try:
+                mod = importlib.import_module(mod_name)
+                setattr(mod, "propagate", faster_propagate)
+            except Exception:
+                pass
+    except ImportError:
+        pass
+
+
+def disable() -> None:
+    """Restore original diffBloch functions."""
+    pass

faster_diffbloch/builder.py

@@ -0,0 +1,58 @@
+"""Native runtime builder and loader for faster-diffbloch."""
+
+from __future__ import annotations
+
+import ctypes
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+NATIVE_DIR = Path(__file__).resolve().parent / "native"
+BUILD_DIR = Path.home() / ".cache" / "faster_diffbloch"
+
+
+def build_and_load_library() -> ctypes.CDLL:
+    """Build or load the cached native acceleration library."""
+    BUILD_DIR.mkdir(parents=True, exist_ok=True)
+    lib_name = "libfaster_diffbloch.dylib" if sys.platform == "darwin" else "libfaster_diffbloch.so"
+    lib_path = BUILD_DIR / lib_name
+
+    sources = [
+        NATIVE_DIR / "bridge_lib.c",
+        NATIVE_DIR / "batch_cgemm.c",
+        NATIVE_DIR / "native_scattering.c",
+    ]
+    if sys.platform == "darwin":
+        sources.append(NATIVE_DIR / "metal_batch_cgemm.m")
+        # Build metallib if metal is available
+        metal_src = NATIVE_DIR / "batch_cgemm.metal"
+        metallib_path = BUILD_DIR / "batch_cgemm.metallib"
+        if metal_src.exists() and (not metallib_path.exists() or metallib_path.stat().st_mtime < metal_src.stat().st_mtime):
+            try:
+                subprocess.run(
+                    ["xcrun", "-sdk", "macosx", "metal", "-O3", "-c", str(metal_src), "-o", str(BUILD_DIR / "batch_cgemm.air")],
+                    check=True, capture_output=True
+                )
+                subprocess.run(
+                    ["xcrun", "-sdk", "macosx", "metallib", str(BUILD_DIR / "batch_cgemm.air"), "-o", str(metallib_path)],
+                    check=True, capture_output=True
+                )
+            except Exception:
+                pass
+
+    stale = not lib_path.exists() or any(lib_path.stat().st_mtime < s.stat().st_mtime for s in sources if s.exists())
+    if stale:
+        frameworks = ["-framework", "Accelerate", "-framework", "Metal", "-framework", "Foundation"] if sys.platform == "darwin" else ["-lblas"]
+        cmd = [
+            "clang", "-std=c11", "-O3", "-fPIC", "-shared",
+            "-D_DEFAULT_SOURCE", "-Wno-unused-function", "-Wno-unused-variable",
+            f"-I{NATIVE_DIR}",
+            *[str(s) for s in sources if s.exists()],
+            *frameworks, "-lm", "-o", str(lib_path),
+        ]
+        res = subprocess.run(cmd, capture_output=True, text=True)
+        if res.returncode != 0:
+            raise RuntimeError(f"Failed to build faster-diffbloch native library:\n{res.stderr}")
+
+    return ctypes.CDLL(str(lib_path))

faster_diffbloch/cli.py

@@ -0,0 +1,21 @@
+"""CLI wrapper enabling faster-diffbloch before running diffbloch commands."""
+
+from __future__ import annotations
+
+import sys
+from .backend import enable
+
+
+def main() -> None:
+    # Enable Metal GPU acceleration by default
+    enable(device="gpu")
+    try:
+        from diffBloch.app.cli import main as diffbloch_main
+        diffbloch_main()
+    except ImportError:
+        print("Error: diffbloch is not installed. Install diffbloch or run 'pip install diffbloch'.", file=sys.stderr)
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()

faster_diffbloch/native/batch_cgemm.c

@@ -0,0 +1,214 @@
+#include "batch_cgemm.h"
+#include <Accelerate/Accelerate.h>
+#include <dispatch/dispatch.h>
+#include <arm_neon.h>
+#include <stdlib.h>
+#include <string.h>
+
+static inline int to_cblas_trans(BatchGemmTranspose t) {
+    switch (t) {
+        case BATCH_CGEMM_NO_TRANS:   return CblasNoTrans;
+        case BATCH_CGEMM_TRANS:      return CblasTrans;
+        case BATCH_CGEMM_CONJ_TRANS: return CblasConjTrans;
+        default:                     return CblasNoTrans;
+    }
+}
+
+// Small matrix microkernel for c64 when M, N, K <= 16 and NoTrans
+static inline void micro_cgemm_c64_small(
+    int32_t M, int32_t N, int32_t K,
+    complex64_t alpha,
+    const complex64_t* A, int32_t lda,
+    const complex64_t* B, int32_t ldb,
+    complex64_t beta,
+    complex64_t* C, int32_t ldc
+) {
+    float ar = crealf(alpha);
+    float ai = cimagf(alpha);
+    float br = crealf(beta);
+    float bi = cimagf(beta);
+
+    for (int i = 0; i < M; i++) {
+        for (int j = 0; j < N; j++) {
+            float acc_r = 0.0f;
+            float acc_i = 0.0f;
+            for (int k = 0; k < K; k++) {
+                complex64_t a_val = A[i * lda + k];
+                complex64_t b_val = B[k * ldb + j];
+                float ak_r = crealf(a_val);
+                float ak_i = cimagf(a_val);
+                float bk_r = crealf(b_val);
+                float bk_i = cimagf(b_val);
+                acc_r += (ak_r * bk_r - ak_i * bk_i);
+                acc_i += (ak_r * bk_i + ak_i * bk_r);
+            }
+            float out_r = acc_r * ar - acc_i * ai;
+            float out_i = acc_r * ai + acc_i * ar;
+            if (br != 0.0f || bi != 0.0f) {
+                complex64_t c_val = C[i * ldc + j];
+                float cr = crealf(c_val);
+                float ci = cimagf(c_val);
+                out_r += (cr * br - ci * bi);
+                out_i += (cr * bi + ci * br);
+            }
+            C[i * ldc + j] = out_r + I * out_i;
+        }
+    }
+}
+
+void batch_cgemm_strided_cpu_c64(
+    BatchGemmTranspose transA,
+    BatchGemmTranspose transB,
+    int32_t M,
+    int32_t N,
+    int32_t K,
+    const complex64_t* alpha,
+    const complex64_t* A,
+    int32_t lda,
+    int64_t strideA,
+    const complex64_t* B,
+    int32_t ldb,
+    int64_t strideB,
+    const complex64_t* beta,
+    complex64_t* C,
+    int32_t ldc,
+    int64_t strideC,
+    int32_t batch_count
+) {
+    if (batch_count <= 0) return;
+
+    int cblas_ta = to_cblas_trans(transA);
+    int cblas_tb = to_cblas_trans(transB);
+
+    if (batch_count == 1) {
+        cblas_cgemm(
+            CblasRowMajor,
+            cblas_ta,
+            cblas_tb,
+            M, N, K,
+            alpha,
+            (const void*)A, lda,
+            (const void*)B, ldb,
+            beta,
+            (void*)C, ldc
+        );
+        return;
+    }
+
+    // Parallelize batch slices across cores using Grand Central Dispatch
+    dispatch_apply(batch_count, dispatch_get_global_queue(DISPATCH_QUEUE_PRIORITY_HIGH, 0), ^(size_t b) {
+        const complex64_t* cur_a = A + b * strideA;
+        const complex64_t* cur_b = B + b * strideB;
+        complex64_t* cur_c = C + b * strideC;
+
+        if (M <= 16 && N <= 16 && K <= 16 && transA == BATCH_CGEMM_NO_TRANS && transB == BATCH_CGEMM_NO_TRANS) {
+            micro_cgemm_c64_small(M, N, K, *alpha, cur_a, lda, cur_b, ldb, *beta, cur_c, ldc);
+        } else {
+            cblas_cgemm(
+                CblasRowMajor,
+                cblas_ta,
+                cblas_tb,
+                M, N, K,
+                alpha,
+                (const void*)cur_a, lda,
+                (const void*)cur_b, ldb,
+                beta,
+                (void*)cur_c, ldc
+            );
+        }
+    });
+}
+
+void batch_cgemm_strided_cpu_c128(
+    BatchGemmTranspose transA,
+    BatchGemmTranspose transB,
+    int32_t M,
+    int32_t N,
+    int32_t K,
+    const complex128_t* alpha,
+    const complex128_t* A,
+    int32_t lda,
+    int64_t strideA,
+    const complex128_t* B,
+    int32_t ldb,
+    int64_t strideB,
+    const complex128_t* beta,
+    complex128_t* C,
+    int32_t ldc,
+    int64_t strideC,
+    int32_t batch_count
+) {
+    if (batch_count <= 0) return;
+
+    int cblas_ta = to_cblas_trans(transA);
+    int cblas_tb = to_cblas_trans(transB);
+
+    if (batch_count == 1) {
+        cblas_zgemm(
+            CblasRowMajor,
+            cblas_ta,
+            cblas_tb,
+            M, N, K,
+            alpha,
+            (const void*)A, lda,
+            (const void*)B, ldb,
+            beta,
+            (void*)C, ldc
+        );
+        return;
+    }
+
+    dispatch_apply(batch_count, dispatch_get_global_queue(DISPATCH_QUEUE_PRIORITY_HIGH, 0), ^(size_t b) {
+        const complex128_t* cur_a = A + b * strideA;
+        const complex128_t* cur_b = B + b * strideB;
+        complex128_t* cur_c = C + b * strideC;
+
+        cblas_zgemm(
+            CblasRowMajor,
+            cblas_ta,
+            cblas_tb,
+            M, N, K,
+            alpha,
+            (const void*)cur_a, lda,
+            (const void*)cur_b, ldb,
+            beta,
+            (void*)cur_c, ldc
+        );
+    });
+}
+
+void batch_cgemm_ptrs_cpu_c64(
+    BatchGemmTranspose transA,
+    BatchGemmTranspose transB,
+    int32_t M,
+    int32_t N,
+    int32_t K,
+    const complex64_t* alpha,
+    const complex64_t* const* A_ptrs,
+    int32_t lda,
+    const complex64_t* const* B_ptrs,
+    int32_t ldb,
+    const complex64_t* beta,
+    complex64_t* const* C_ptrs,
+    int32_t ldc,
+    int32_t batch_count
+) {
+    if (batch_count <= 0) return;
+
+    int cblas_ta = to_cblas_trans(transA);
+    int cblas_tb = to_cblas_trans(transB);
+
+    dispatch_apply(batch_count, dispatch_get_global_queue(DISPATCH_QUEUE_PRIORITY_HIGH, 0), ^(size_t b) {
+        cblas_cgemm(
+            CblasRowMajor,
+            cblas_ta,
+            cblas_tb,
+            M, N, K,
+            alpha,
+            (const void*)A_ptrs[b], lda,
+            (const void*)B_ptrs[b], ldb,
+            beta,
+            (void*)C_ptrs[b], ldc
+        );
+    });
+}

faster_diffbloch/native/batch_cgemm.h

@@ -0,0 +1,142 @@
+#ifndef BATCH_CGEMM_H
+#define BATCH_CGEMM_H
+
+#include <stdint.h>
+#include <stdbool.h>
+#include <complex.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+typedef float _Complex complex64_t;
+typedef double _Complex complex128_t;
+
+typedef enum {
+    BATCH_CGEMM_ROW_MAJOR = 101,
+    BATCH_CGEMM_COL_MAJOR = 102
+} BatchGemmOrder;
+
+typedef enum {
+    BATCH_CGEMM_NO_TRANS   = 111,
+    BATCH_CGEMM_TRANS      = 112,
+    BATCH_CGEMM_CONJ_TRANS = 113
+} BatchGemmTranspose;
+
+/* CPU Multithreaded Strided Batched Complex GEMM (Float32 / c64) */
+void batch_cgemm_strided_cpu_c64(
+    BatchGemmTranspose transA,
+    BatchGemmTranspose transB,
+    int32_t M,
+    int32_t N,
+    int32_t K,
+    const complex64_t* alpha,
+    const complex64_t* A,
+    int32_t lda,
+    int64_t strideA,
+    const complex64_t* B,
+    int32_t ldb,
+    int64_t strideB,
+    const complex64_t* beta,
+    complex64_t* C,
+    int32_t ldc,
+    int64_t strideC,
+    int32_t batch_count
+);
+
+/* CPU Multithreaded Strided Batched Complex GEMM (Float64 / c128) */
+void batch_cgemm_strided_cpu_c128(
+    BatchGemmTranspose transA,
+    BatchGemmTranspose transB,
+    int32_t M,
+    int32_t N,
+    int32_t K,
+    const complex128_t* alpha,
+    const complex128_t* A,
+    int32_t lda,
+    int64_t strideA,
+    const complex128_t* B,
+    int32_t ldb,
+    int64_t strideB,
+    const complex128_t* beta,
+    complex128_t* C,
+    int32_t ldc,
+    int64_t strideC,
+    int32_t batch_count
+);
+
+/* GPU Metal Strided Batched Complex GEMM (Float32 / c64) */
+int32_t batch_cgemm_strided_metal_c64(
+    BatchGemmTranspose transA,
+    BatchGemmTranspose transB,
+    int32_t M,
+    int32_t N,
+    int32_t K,
+    const complex64_t* alpha,
+    const complex64_t* A,
+    int32_t lda,
+    int64_t strideA,
+    const complex64_t* B,
+    int32_t ldb,
+    int64_t strideB,
+    const complex64_t* beta,
+    complex64_t* C,
+    int32_t ldc,
+    int64_t strideC,
+    int32_t batch_count
+);
+
+/* GPU Metal Batched Matrix Exponential (Float32 / c64) */
+int32_t metal_matrix_exp_c64(
+    const complex64_t* A,
+    int32_t n,
+    int32_t batch_count,
+    complex64_t* out
+);
+
+/* GPU Metal Batched Matrix Exponential Adjoint (Float32 / c64) */
+int32_t metal_matrix_exp_backward_c64(
+    const complex64_t* M,
+    const complex64_t* Ebar,
+    int32_t n,
+    int32_t batch_count,
+    int32_t dense,
+    complex64_t* Mbar
+);
+
+/* GPU Metal Structure Factors (Float64 / c128) */
+int32_t metal_structure_factors_c128(
+    complex128_t* fgb,
+    const double* g2,
+    const double* hkl,
+    const double* lobato_a,
+    const double* lobato_b,
+    const double* uij,
+    const double* pos,
+    const double* occ,
+    int32_t n_grid,
+    int32_t n_atoms,
+    double volume
+);
+
+/* GPU Metal Structure Matrix Assembly (Float64 / c128) */
+int32_t metal_structure_matrix_c128(
+    const complex128_t* fgb,
+    int32_t n_grid,
+    const int32_t* source,
+    int32_t buffer_size,
+    const int32_t* destination,
+    int32_t n_beams,
+    const double* mii,
+    const double* diagonal,
+    int32_t n_batch,
+    double prefactor,
+    int32_t absorption,
+    complex128_t* out
+);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif /* BATCH_CGEMM_H */