faster-diffbloch 0.1.1__tar.gz → 0.1.2__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {faster_diffbloch-0.1.1 → faster_diffbloch-0.1.2}/PKG-INFO +43 -12
- faster_diffbloch-0.1.2/README.md +112 -0
- {faster_diffbloch-0.1.1 → faster_diffbloch-0.1.2}/pyproject.toml +5 -1
- {faster_diffbloch-0.1.1 → faster_diffbloch-0.1.2}/src/faster_diffbloch/__init__.py +1 -1
- {faster_diffbloch-0.1.1 → faster_diffbloch-0.1.2}/src/faster_diffbloch/backend.py +15 -12
- faster_diffbloch-0.1.1/README.md +0 -85
- {faster_diffbloch-0.1.1 → faster_diffbloch-0.1.2}/.gitignore +0 -0
- {faster_diffbloch-0.1.1 → faster_diffbloch-0.1.2}/LICENSE +0 -0
- {faster_diffbloch-0.1.1 → faster_diffbloch-0.1.2}/src/faster_diffbloch/builder.py +0 -0
- {faster_diffbloch-0.1.1 → faster_diffbloch-0.1.2}/src/faster_diffbloch/cli.py +0 -0
- {faster_diffbloch-0.1.1 → faster_diffbloch-0.1.2}/src/faster_diffbloch/native/batch_cgemm.c +0 -0
- {faster_diffbloch-0.1.1 → faster_diffbloch-0.1.2}/src/faster_diffbloch/native/batch_cgemm.h +0 -0
- {faster_diffbloch-0.1.1 → faster_diffbloch-0.1.2}/src/faster_diffbloch/native/batch_cgemm.metal +0 -0
- {faster_diffbloch-0.1.1 → faster_diffbloch-0.1.2}/src/faster_diffbloch/native/bridge_lib.c +0 -0
- {faster_diffbloch-0.1.1 → faster_diffbloch-0.1.2}/src/faster_diffbloch/native/metal_batch_cgemm.m +0 -0
- {faster_diffbloch-0.1.1 → faster_diffbloch-0.1.2}/src/faster_diffbloch/native/native_scattering.c +0 -0
- {faster_diffbloch-0.1.1 → faster_diffbloch-0.1.2}/src/faster_diffbloch/native/native_scattering.h +0 -0
|
@@ -1,14 +1,18 @@
|
|
|
1
1
|
Metadata-Version: 2.5
|
|
2
2
|
Name: faster-diffbloch
|
|
3
|
-
Version: 0.1.
|
|
3
|
+
Version: 0.1.2
|
|
4
4
|
Summary: Drop-in Metal GPU (macOS) and optimized CPU (macOS/Linux) acceleration for diffBloch
|
|
5
5
|
Project-URL: Homepage, https://godofecht.github.io/diffFlow/
|
|
6
|
+
Project-URL: Documentation, https://godofecht.github.io/diffFlow/
|
|
6
7
|
Project-URL: Repository, https://github.com/godofecht/diffFlow
|
|
7
8
|
Project-URL: Issues, https://github.com/godofecht/diffFlow/issues
|
|
9
|
+
Project-URL: PyPI, https://pypi.org/project/faster-diffbloch/
|
|
10
|
+
Project-URL: piwheels, https://www.piwheels.org/project/faster-diffbloch/
|
|
8
11
|
Project-URL: Original diffBloch, https://diffbloch.com
|
|
9
12
|
Author-email: Abhishek Shivakumar <abhishek@example.com>
|
|
10
13
|
License-Expression: MIT
|
|
11
14
|
License-File: LICENSE
|
|
15
|
+
Keywords: Metal GPU,PyTorch acceleration,diffBloch,electron crystallography,electron diffraction,structure refinement
|
|
12
16
|
Classifier: Development Status :: 4 - Beta
|
|
13
17
|
Classifier: Intended Audience :: Science/Research
|
|
14
18
|
Classifier: License :: OSI Approved :: MIT License
|
|
@@ -27,13 +31,15 @@ Provides-Extra: diffbloch
|
|
|
27
31
|
Requires-Dist: diffbloch; extra == 'diffbloch'
|
|
28
32
|
Description-Content-Type: text/markdown
|
|
29
33
|
|
|
30
|
-
# faster-
|
|
34
|
+
# faster-diffbloch
|
|
31
35
|
|
|
32
36
|
Drop-in Apple Silicon Metal GPU and optimized CPU acceleration for [diffBloch](https://diffbloch.com) electron crystallography structure refinement.
|
|
33
37
|
|
|
34
|
-
|
|
38
|
+
**Up to 1.80× faster than PyTorch CPU** in measured end-to-end forward-plus-backward benchmarks, with a Metal GPU path for Apple Silicon and an optimized CPU path for macOS and Linux.
|
|
35
39
|
|
|
36
|
-
Original
|
|
40
|
+
| Package | Documentation | Source repository | Original project |
|
|
41
|
+
| :--- | :--- | :--- | :--- |
|
|
42
|
+
| [PyPI](https://pypi.org/project/faster-diffbloch/) · [piwheels](https://www.piwheels.org/project/faster-diffbloch/) | [diffFlow documentation and benchmarks](https://godofecht.github.io/diffFlow/) | [github.com/godofecht/diffFlow](https://github.com/godofecht/diffFlow) | [diffbloch.com](https://diffbloch.com) |
|
|
37
43
|
|
|
38
44
|
---
|
|
39
45
|
|
|
@@ -52,20 +58,45 @@ Runtime platform guards automatically detect your operating system and hardware
|
|
|
52
58
|
|
|
53
59
|
## Why faster-diffBloch?
|
|
54
60
|
|
|
55
|
-
|
|
56
|
-
|
|
61
|
+
`faster-diffbloch` keeps the diffBloch workflow and public API intact while moving
|
|
62
|
+
the expensive propagation and gradient work onto an accelerated native backend.
|
|
63
|
+
You get the same scientific calculation, with a faster path on supported hardware
|
|
64
|
+
and a safe PyTorch fallback when the native backend is unavailable.
|
|
57
65
|
|
|
58
|
-
|
|
59
|
-
|
|
60
|
-
|
|
61
|
-
3. **Bit-for-Bit Validation:**
|
|
62
|
-
Passes all 738 unit tests in diffBloch and reproduces the experimental 99-rotation quartz dataset ($R_{\text{obs}} = 0.0486$).
|
|
66
|
+
The package is validated against the diffBloch test suite and reference results.
|
|
67
|
+
It passes all 738 diffBloch tests, 55 additional conformance tests, and reproduces
|
|
68
|
+
the experimental quartz refinement result.
|
|
63
69
|
|
|
64
70
|
---
|
|
65
71
|
|
|
66
72
|
## Performance
|
|
67
73
|
|
|
68
|
-
|
|
74
|
+
End-to-end forward-plus-backward timing on an Apple M4 Max, measured across the
|
|
75
|
+
same diffBloch workload at different beam counts:
|
|
76
|
+
|
|
77
|
+
| Beams | faster-diffBloch | PyTorch | Speedup |
|
|
78
|
+
| :---: | :---: | :---: | :---: |
|
|
79
|
+
| 31 | 0.90 ms | 1.62 ms | **1.80×** |
|
|
80
|
+
| 61 | 1.37 ms | 2.24 ms | **1.63×** |
|
|
81
|
+
| 91 | 2.84 ms | 3.99 ms | **1.41×** |
|
|
82
|
+
| 163 | 7.67 ms | 10.98 ms | **1.43×** |
|
|
83
|
+
| 579 | 89.81 ms | 135.77 ms | **1.51×** |
|
|
84
|
+
|
|
85
|
+
That is **1.41×–1.80× faster than PyTorch CPU** across the measured cases,
|
|
86
|
+
including the largest 579-beam case. The same benchmark is **1.21× faster than
|
|
87
|
+
the Mojo reference at 31 beams**, remains ahead through 163 beams, and is
|
|
88
|
+
essentially level at 579 beams.
|
|
89
|
+
|
|
90
|
+
On Apple Silicon, the Metal GPU path is also available. In the package's M4 Max
|
|
91
|
+
comparison run at 579 beams, it measured **2.24× faster than PyTorch CPU for the
|
|
92
|
+
forward pass** and **2.63× faster than PyTorch MPS for forward plus backward**.
|
|
93
|
+
|
|
94
|
+
These figures are workload benchmarks, not a promise that every crystal,
|
|
95
|
+
hardware configuration, or beam count will see the same result. The CPU table
|
|
96
|
+
uses single-threaded PyTorch and the same Apple Accelerate environment for a
|
|
97
|
+
like-for-like comparison.
|
|
98
|
+
|
|
99
|
+
### Package benchmark snapshot
|
|
69
100
|
|
|
70
101
|
| Implementation | Forward | Forward + Backward | Speedup vs PyTorch CPU | Speedup vs PyTorch MPS |
|
|
71
102
|
| :--- | :---: | :---: | :---: | :---: |
|
|
@@ -0,0 +1,112 @@
|
|
|
1
|
+
# faster-diffbloch
|
|
2
|
+
|
|
3
|
+
Drop-in Apple Silicon Metal GPU and optimized CPU acceleration for [diffBloch](https://diffbloch.com) electron crystallography structure refinement.
|
|
4
|
+
|
|
5
|
+
**Up to 1.80× faster than PyTorch CPU** in measured end-to-end forward-plus-backward benchmarks, with a Metal GPU path for Apple Silicon and an optimized CPU path for macOS and Linux.
|
|
6
|
+
|
|
7
|
+
| Package | Documentation | Source repository | Original project |
|
|
8
|
+
| :--- | :--- | :--- | :--- |
|
|
9
|
+
| [PyPI](https://pypi.org/project/faster-diffbloch/) · [piwheels](https://www.piwheels.org/project/faster-diffbloch/) | [diffFlow documentation and benchmarks](https://godofecht.github.io/diffFlow/) | [github.com/godofecht/diffFlow](https://github.com/godofecht/diffFlow) | [diffbloch.com](https://diffbloch.com) |
|
|
10
|
+
|
|
11
|
+
---
|
|
12
|
+
|
|
13
|
+
## Operating System and Platform Support
|
|
14
|
+
|
|
15
|
+
| Operating System / Hardware | CPU Acceleration (`device="cpu"`) | GPU Acceleration (`device="gpu"`) | Backend Runtime |
|
|
16
|
+
| :--- | :---: | :---: | :--- |
|
|
17
|
+
| **macOS Apple Silicon (M1/M2/M3/M4/Max/Ultra)** | **Supported** | **Supported** | Native Metal Compute Shaders + Apple Accelerate BLAS |
|
|
18
|
+
| **macOS Intel (x86_64)** | **Supported** | Fallback to CPU | Apple Accelerate BLAS |
|
|
19
|
+
| **Linux (x86_64 / aarch64)** | **Supported** | Fallback to CPU | OpenBLAS / C11 BLAS |
|
|
20
|
+
| **Windows** | PyTorch Reference | PyTorch Reference | Pure PyTorch Reference Fallback |
|
|
21
|
+
|
|
22
|
+
Runtime platform guards automatically detect your operating system and hardware configuration. When `device="gpu"` is requested on Linux, `faster-diffbloch` automatically selects the optimized CPU backend with an informative warning.
|
|
23
|
+
|
|
24
|
+
---
|
|
25
|
+
|
|
26
|
+
## Why faster-diffBloch?
|
|
27
|
+
|
|
28
|
+
`faster-diffbloch` keeps the diffBloch workflow and public API intact while moving
|
|
29
|
+
the expensive propagation and gradient work onto an accelerated native backend.
|
|
30
|
+
You get the same scientific calculation, with a faster path on supported hardware
|
|
31
|
+
and a safe PyTorch fallback when the native backend is unavailable.
|
|
32
|
+
|
|
33
|
+
The package is validated against the diffBloch test suite and reference results.
|
|
34
|
+
It passes all 738 diffBloch tests, 55 additional conformance tests, and reproduces
|
|
35
|
+
the experimental quartz refinement result.
|
|
36
|
+
|
|
37
|
+
---
|
|
38
|
+
|
|
39
|
+
## Performance
|
|
40
|
+
|
|
41
|
+
End-to-end forward-plus-backward timing on an Apple M4 Max, measured across the
|
|
42
|
+
same diffBloch workload at different beam counts:
|
|
43
|
+
|
|
44
|
+
| Beams | faster-diffBloch | PyTorch | Speedup |
|
|
45
|
+
| :---: | :---: | :---: | :---: |
|
|
46
|
+
| 31 | 0.90 ms | 1.62 ms | **1.80×** |
|
|
47
|
+
| 61 | 1.37 ms | 2.24 ms | **1.63×** |
|
|
48
|
+
| 91 | 2.84 ms | 3.99 ms | **1.41×** |
|
|
49
|
+
| 163 | 7.67 ms | 10.98 ms | **1.43×** |
|
|
50
|
+
| 579 | 89.81 ms | 135.77 ms | **1.51×** |
|
|
51
|
+
|
|
52
|
+
That is **1.41×–1.80× faster than PyTorch CPU** across the measured cases,
|
|
53
|
+
including the largest 579-beam case. The same benchmark is **1.21× faster than
|
|
54
|
+
the Mojo reference at 31 beams**, remains ahead through 163 beams, and is
|
|
55
|
+
essentially level at 579 beams.
|
|
56
|
+
|
|
57
|
+
On Apple Silicon, the Metal GPU path is also available. In the package's M4 Max
|
|
58
|
+
comparison run at 579 beams, it measured **2.24× faster than PyTorch CPU for the
|
|
59
|
+
forward pass** and **2.63× faster than PyTorch MPS for forward plus backward**.
|
|
60
|
+
|
|
61
|
+
These figures are workload benchmarks, not a promise that every crystal,
|
|
62
|
+
hardware configuration, or beam count will see the same result. The CPU table
|
|
63
|
+
uses single-threaded PyTorch and the same Apple Accelerate environment for a
|
|
64
|
+
like-for-like comparison.
|
|
65
|
+
|
|
66
|
+
### Package benchmark snapshot
|
|
67
|
+
|
|
68
|
+
| Implementation | Forward | Forward + Backward | Speedup vs PyTorch CPU | Speedup vs PyTorch MPS |
|
|
69
|
+
| :--- | :---: | :---: | :---: | :---: |
|
|
70
|
+
| PyTorch CPU | 25.7 ms | 130.3 ms | 1.00x | 1.17x |
|
|
71
|
+
| PyTorch MPS (fallback) | 26.2 ms | 153.0 ms | 0.85x | 1.00x |
|
|
72
|
+
| **faster-diffBloch CPU** | **24.4 ms** | **83.5 ms** | **1.56x** | **1.83x** |
|
|
73
|
+
| **faster-diffBloch Metal GPU** | **13.1 ms** | **58.1 ms** | **2.24x** | **2.63x** |
|
|
74
|
+
|
|
75
|
+
---
|
|
76
|
+
|
|
77
|
+
## Installation
|
|
78
|
+
|
|
79
|
+
```bash
|
|
80
|
+
pip install faster-diffbloch
|
|
81
|
+
```
|
|
82
|
+
|
|
83
|
+
---
|
|
84
|
+
|
|
85
|
+
## Usage
|
|
86
|
+
|
|
87
|
+
### 1. Drop-in CLI
|
|
88
|
+
|
|
89
|
+
Use `diffbloch-fast` or `faster-diffbloch` anywhere you would use `diffbloch`:
|
|
90
|
+
|
|
91
|
+
```bash
|
|
92
|
+
diffbloch-fast infer examples/Colmey_et_al_2026/data/quartz-no-abs
|
|
93
|
+
diffbloch-fast refine examples/Colmey_et_al_2026/data/quartz-no-abs
|
|
94
|
+
```
|
|
95
|
+
|
|
96
|
+
### 2. Python API Injection
|
|
97
|
+
|
|
98
|
+
Enable acceleration inside any existing diffBloch script:
|
|
99
|
+
|
|
100
|
+
```python
|
|
101
|
+
import faster_diffbloch
|
|
102
|
+
|
|
103
|
+
# Enable Metal GPU acceleration (macOS Apple Silicon)
|
|
104
|
+
faster_diffbloch.enable(device="gpu")
|
|
105
|
+
|
|
106
|
+
# Or CPU acceleration (macOS and Linux)
|
|
107
|
+
faster_diffbloch.enable(device="cpu")
|
|
108
|
+
|
|
109
|
+
# Run standard diffBloch code
|
|
110
|
+
import diffBloch
|
|
111
|
+
# All propagate and matrix_exp calls now route through faster-diffBloch
|
|
112
|
+
```
|
|
@@ -4,10 +4,11 @@ build-backend = "hatchling.build"
|
|
|
4
4
|
|
|
5
5
|
[project]
|
|
6
6
|
name = "faster-diffbloch"
|
|
7
|
-
version = "0.1.
|
|
7
|
+
version = "0.1.2"
|
|
8
8
|
description = "Drop-in Metal GPU (macOS) and optimized CPU (macOS/Linux) acceleration for diffBloch"
|
|
9
9
|
readme = "README.md"
|
|
10
10
|
requires-python = ">=3.10"
|
|
11
|
+
keywords = ["diffBloch", "electron diffraction", "electron crystallography", "structure refinement", "PyTorch acceleration", "Metal GPU"]
|
|
11
12
|
license = "MIT"
|
|
12
13
|
authors = [
|
|
13
14
|
{ name = "Abhishek Shivakumar", email = "abhishek@example.com" }
|
|
@@ -41,8 +42,11 @@ diffbloch-fast = "faster_diffbloch.cli:main"
|
|
|
41
42
|
|
|
42
43
|
[project.urls]
|
|
43
44
|
Homepage = "https://godofecht.github.io/diffFlow/"
|
|
45
|
+
Documentation = "https://godofecht.github.io/diffFlow/"
|
|
44
46
|
Repository = "https://github.com/godofecht/diffFlow"
|
|
45
47
|
Issues = "https://github.com/godofecht/diffFlow/issues"
|
|
48
|
+
PyPI = "https://pypi.org/project/faster-diffbloch/"
|
|
49
|
+
piwheels = "https://www.piwheels.org/project/faster-diffbloch/"
|
|
46
50
|
"Original diffBloch" = "https://diffbloch.com"
|
|
47
51
|
|
|
48
52
|
[tool.hatch.build.targets.wheel]
|
|
@@ -46,15 +46,15 @@ def matrix_exp(a: np.ndarray, device: str = "gpu") -> np.ndarray:
|
|
|
46
46
|
fn = getattr(lib, "bridge_matrix_exp_gpu_ptr_c64_i32_i32_ptr_c64", None)
|
|
47
47
|
if fn:
|
|
48
48
|
fn.argtypes = [ctypes.c_void_p, ctypes.c_int32, ctypes.c_int32, ctypes.c_void_p]
|
|
49
|
-
fn.restype =
|
|
50
|
-
fn(ptr_in,
|
|
51
|
-
|
|
49
|
+
fn.restype = ctypes.c_int32
|
|
50
|
+
if fn(ptr_in, n, batch, ptr_out) == 0:
|
|
51
|
+
return out
|
|
52
52
|
|
|
53
53
|
fn = getattr(lib, "bridge_matrix_exp_ptr_c64_i32_i32_ptr_c64", None)
|
|
54
54
|
if fn:
|
|
55
55
|
fn.argtypes = [ctypes.c_void_p, ctypes.c_int32, ctypes.c_int32, ctypes.c_void_p]
|
|
56
56
|
fn.restype = None
|
|
57
|
-
fn(ptr_in,
|
|
57
|
+
fn(ptr_in, n, batch, ptr_out)
|
|
58
58
|
return out
|
|
59
59
|
|
|
60
60
|
return torch.matrix_exp(torch.from_numpy(a)).numpy()
|
|
@@ -63,10 +63,13 @@ def matrix_exp(a: np.ndarray, device: str = "gpu") -> np.ndarray:
|
|
|
63
63
|
def matrix_exp_backward(a: np.ndarray, ebar: np.ndarray, dense: bool = False, device: str = "gpu") -> np.ndarray:
|
|
64
64
|
lib = get_library()
|
|
65
65
|
if lib is None:
|
|
66
|
-
|
|
67
|
-
|
|
68
|
-
|
|
69
|
-
|
|
66
|
+
# Custom autograd.Function backward methods run with grad mode disabled.
|
|
67
|
+
# Re-enable it for the correctness fallback.
|
|
68
|
+
with torch.enable_grad():
|
|
69
|
+
tA = torch.from_numpy(a).requires_grad_(True)
|
|
70
|
+
tE = torch.matrix_exp(tA)
|
|
71
|
+
tE.backward(torch.from_numpy(ebar))
|
|
72
|
+
return tA.grad.numpy()
|
|
70
73
|
|
|
71
74
|
a_c = np.ascontiguousarray(a, dtype=np.complex64)
|
|
72
75
|
ebar_c = np.ascontiguousarray(ebar, dtype=np.complex64)
|
|
@@ -83,15 +86,15 @@ def matrix_exp_backward(a: np.ndarray, ebar: np.ndarray, dense: bool = False, de
|
|
|
83
86
|
fn = getattr(lib, "bridge_matrix_exp_backward_gpu_ptr_c64_ptr_c64_i32_i32_i32_ptr_c64", None)
|
|
84
87
|
if fn:
|
|
85
88
|
fn.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_int32, ctypes.c_int32, ctypes.c_int32, ctypes.c_void_p]
|
|
86
|
-
fn.restype =
|
|
87
|
-
fn(ptr_a, ptr_ebar,
|
|
88
|
-
|
|
89
|
+
fn.restype = ctypes.c_int32
|
|
90
|
+
if fn(ptr_a, ptr_ebar, n, batch, dense_val, ptr_out) == 0:
|
|
91
|
+
return out
|
|
89
92
|
|
|
90
93
|
fn = getattr(lib, "bridge_matrix_exp_backward_ptr_c64_ptr_c64_i32_i32_i32_ptr_c64", None)
|
|
91
94
|
if fn:
|
|
92
95
|
fn.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_int32, ctypes.c_int32, ctypes.c_int32, ctypes.c_void_p]
|
|
93
96
|
fn.restype = None
|
|
94
|
-
fn(ptr_a, ptr_ebar,
|
|
97
|
+
fn(ptr_a, ptr_ebar, n, batch, dense_val, ptr_out)
|
|
95
98
|
return out
|
|
96
99
|
|
|
97
100
|
tA = torch.from_numpy(a).requires_grad_(True)
|
faster_diffbloch-0.1.1/README.md
DELETED
|
@@ -1,85 +0,0 @@
|
|
|
1
|
-
# faster-diffBloch
|
|
2
|
-
|
|
3
|
-
Drop-in Apple Silicon Metal GPU and optimized CPU acceleration for [diffBloch](https://diffbloch.com) electron crystallography structure refinement.
|
|
4
|
-
|
|
5
|
-
Documentation and comparison benchmarks: [https://godofecht.github.io/diffFlow/](https://godofecht.github.io/diffFlow/)
|
|
6
|
-
|
|
7
|
-
Original diffBloch project: [https://diffbloch.com](https://diffbloch.com)
|
|
8
|
-
|
|
9
|
-
---
|
|
10
|
-
|
|
11
|
-
## Operating System and Platform Support
|
|
12
|
-
|
|
13
|
-
| Operating System / Hardware | CPU Acceleration (`device="cpu"`) | GPU Acceleration (`device="gpu"`) | Backend Runtime |
|
|
14
|
-
| :--- | :---: | :---: | :--- |
|
|
15
|
-
| **macOS Apple Silicon (M1/M2/M3/M4/Max/Ultra)** | **Supported** | **Supported** | Native Metal Compute Shaders + Apple Accelerate BLAS |
|
|
16
|
-
| **macOS Intel (x86_64)** | **Supported** | Fallback to CPU | Apple Accelerate BLAS |
|
|
17
|
-
| **Linux (x86_64 / aarch64)** | **Supported** | Fallback to CPU | OpenBLAS / C11 BLAS |
|
|
18
|
-
| **Windows** | PyTorch Reference | PyTorch Reference | Pure PyTorch Reference Fallback |
|
|
19
|
-
|
|
20
|
-
Runtime platform guards automatically detect your operating system and hardware configuration. When `device="gpu"` is requested on Linux, `faster-diffbloch` automatically selects the optimized CPU backend with an informative warning.
|
|
21
|
-
|
|
22
|
-
---
|
|
23
|
-
|
|
24
|
-
## Why faster-diffBloch?
|
|
25
|
-
|
|
26
|
-
1. **Native Metal GPU Execution on macOS:**
|
|
27
|
-
PyTorch MPS lacks a native GPU kernel for `aten::linalg_matrix_exp`, which causes PyTorch to fall back to CPU execution with host-device memory transfers. `faster-diffBloch` executes matrix exponentials directly on Apple Silicon Metal with zero-copy unified memory.
|
|
28
|
-
|
|
29
|
-
2. **Blocked-Pair Adjoint Formulation:**
|
|
30
|
-
Standard matrix exponential autograd embeds the operator into a $2N \times 2N$ block matrix, costing $8 \times N^3$ FLOPs. `faster-diffBloch` evaluates the pullback in the block-triangular pair algebra $(Y_a Y_b, Y_a L_b + L_a Y_b)$, reducing the work to $3 \times N^3$ FLOPs (2.67x fewer products).
|
|
31
|
-
|
|
32
|
-
3. **Bit-for-Bit Validation:**
|
|
33
|
-
Passes all 738 unit tests in diffBloch and reproduces the experimental 99-rotation quartz dataset ($R_{\text{obs}} = 0.0486$).
|
|
34
|
-
|
|
35
|
-
---
|
|
36
|
-
|
|
37
|
-
## Performance
|
|
38
|
-
|
|
39
|
-
Forward and backward timing comparison on Apple Silicon (M4 Max) at $N=579$ beams (CsPbBr3 scale):
|
|
40
|
-
|
|
41
|
-
| Implementation | Forward | Forward + Backward | Speedup vs PyTorch CPU | Speedup vs PyTorch MPS |
|
|
42
|
-
| :--- | :---: | :---: | :---: | :---: |
|
|
43
|
-
| PyTorch CPU | 25.7 ms | 130.3 ms | 1.00x | 1.17x |
|
|
44
|
-
| PyTorch MPS (fallback) | 26.2 ms | 153.0 ms | 0.85x | 1.00x |
|
|
45
|
-
| **faster-diffBloch CPU** | **24.4 ms** | **83.5 ms** | **1.56x** | **1.83x** |
|
|
46
|
-
| **faster-diffBloch Metal GPU** | **13.1 ms** | **58.1 ms** | **2.24x** | **2.63x** |
|
|
47
|
-
|
|
48
|
-
---
|
|
49
|
-
|
|
50
|
-
## Installation
|
|
51
|
-
|
|
52
|
-
```bash
|
|
53
|
-
pip install faster-diffbloch
|
|
54
|
-
```
|
|
55
|
-
|
|
56
|
-
---
|
|
57
|
-
|
|
58
|
-
## Usage
|
|
59
|
-
|
|
60
|
-
### 1. Drop-in CLI
|
|
61
|
-
|
|
62
|
-
Use `diffbloch-fast` or `faster-diffbloch` anywhere you would use `diffbloch`:
|
|
63
|
-
|
|
64
|
-
```bash
|
|
65
|
-
diffbloch-fast infer examples/Colmey_et_al_2026/data/quartz-no-abs
|
|
66
|
-
diffbloch-fast refine examples/Colmey_et_al_2026/data/quartz-no-abs
|
|
67
|
-
```
|
|
68
|
-
|
|
69
|
-
### 2. Python API Injection
|
|
70
|
-
|
|
71
|
-
Enable acceleration inside any existing diffBloch script:
|
|
72
|
-
|
|
73
|
-
```python
|
|
74
|
-
import faster_diffbloch
|
|
75
|
-
|
|
76
|
-
# Enable Metal GPU acceleration (macOS Apple Silicon)
|
|
77
|
-
faster_diffbloch.enable(device="gpu")
|
|
78
|
-
|
|
79
|
-
# Or CPU acceleration (macOS and Linux)
|
|
80
|
-
faster_diffbloch.enable(device="cpu")
|
|
81
|
-
|
|
82
|
-
# Run standard diffBloch code
|
|
83
|
-
import diffBloch
|
|
84
|
-
# All propagate and matrix_exp calls now route through faster-diffBloch
|
|
85
|
-
```
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
{faster_diffbloch-0.1.1 → faster_diffbloch-0.1.2}/src/faster_diffbloch/native/batch_cgemm.metal
RENAMED
|
File without changes
|
|
File without changes
|
{faster_diffbloch-0.1.1 → faster_diffbloch-0.1.2}/src/faster_diffbloch/native/metal_batch_cgemm.m
RENAMED
|
File without changes
|
{faster_diffbloch-0.1.1 → faster_diffbloch-0.1.2}/src/faster_diffbloch/native/native_scattering.c
RENAMED
|
File without changes
|
{faster_diffbloch-0.1.1 → faster_diffbloch-0.1.2}/src/faster_diffbloch/native/native_scattering.h
RENAMED
|
File without changes
|