trnblas 0.2.0__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.
Files changed (42) hide show
  1. trnblas-0.2.0/.github/ISSUE_TEMPLATE/bug_report.md +31 -0
  2. trnblas-0.2.0/.github/ISSUE_TEMPLATE/feature_request.md +17 -0
  3. trnblas-0.2.0/.github/pull_request_template.md +9 -0
  4. trnblas-0.2.0/.github/workflows/ci.yml +20 -0
  5. trnblas-0.2.0/.github/workflows/docs.yml +38 -0
  6. trnblas-0.2.0/.github/workflows/neuron.yml +50 -0
  7. trnblas-0.2.0/.github/workflows/publish.yml +56 -0
  8. trnblas-0.2.0/.gitignore +18 -0
  9. trnblas-0.2.0/CHANGELOG.md +46 -0
  10. trnblas-0.2.0/CLAUDE.md +105 -0
  11. trnblas-0.2.0/LICENSE +17 -0
  12. trnblas-0.2.0/PKG-INFO +139 -0
  13. trnblas-0.2.0/README.md +111 -0
  14. trnblas-0.2.0/benchmarks/bench_blas.py +38 -0
  15. trnblas-0.2.0/benchmarks/conftest.py +16 -0
  16. trnblas-0.2.0/docs/api/level1.md +28 -0
  17. trnblas-0.2.0/docs/api/level2.md +20 -0
  18. trnblas-0.2.0/docs/api/level3.md +32 -0
  19. trnblas-0.2.0/docs/api/nki.md +29 -0
  20. trnblas-0.2.0/docs/architecture.md +52 -0
  21. trnblas-0.2.0/docs/index.md +31 -0
  22. trnblas-0.2.0/docs/installation.md +33 -0
  23. trnblas-0.2.0/docs/quickstart.md +53 -0
  24. trnblas-0.2.0/examples/df_mp2.py +147 -0
  25. trnblas-0.2.0/mkdocs.yml +32 -0
  26. trnblas-0.2.0/pyproject.toml +50 -0
  27. trnblas-0.2.0/setup.cfg +4 -0
  28. trnblas-0.2.0/tests/conftest.py +42 -0
  29. trnblas-0.2.0/tests/test_level1.py +62 -0
  30. trnblas-0.2.0/tests/test_level2.py +78 -0
  31. trnblas-0.2.0/tests/test_level3.py +151 -0
  32. trnblas-0.2.0/trnblas/__init__.py +35 -0
  33. trnblas-0.2.0/trnblas/level1.py +45 -0
  34. trnblas-0.2.0/trnblas/level2.py +86 -0
  35. trnblas-0.2.0/trnblas/level3.py +203 -0
  36. trnblas-0.2.0/trnblas/nki/__init__.py +5 -0
  37. trnblas-0.2.0/trnblas/nki/dispatch.py +121 -0
  38. trnblas-0.2.0/trnblas.egg-info/PKG-INFO +139 -0
  39. trnblas-0.2.0/trnblas.egg-info/SOURCES.txt +40 -0
  40. trnblas-0.2.0/trnblas.egg-info/dependency_links.txt +1 -0
  41. trnblas-0.2.0/trnblas.egg-info/requires.txt +11 -0
  42. trnblas-0.2.0/trnblas.egg-info/top_level.txt +1 -0
@@ -0,0 +1,31 @@
1
+ ---
2
+ name: Bug report
3
+ about: Report a bug in trnblas
4
+ labels: bug
5
+ ---
6
+
7
+ ## Environment
8
+
9
+ - trnblas version:
10
+ - Python version:
11
+ - PyTorch version:
12
+ - Hardware: CPU / trn1 / trn2 / inf2
13
+ - OS:
14
+
15
+ ## Description
16
+
17
+ What happened?
18
+
19
+ ## Steps to reproduce
20
+
21
+ ```python
22
+ # Minimal code to reproduce
23
+ ```
24
+
25
+ ## Expected behavior
26
+
27
+ What should have happened?
28
+
29
+ ## Actual behavior
30
+
31
+ What happened instead? Include error messages/tracebacks if applicable.
@@ -0,0 +1,17 @@
1
+ ---
2
+ name: Feature request
3
+ about: Suggest a new feature or enhancement
4
+ labels: enhancement
5
+ ---
6
+
7
+ ## Description
8
+
9
+ What feature would you like?
10
+
11
+ ## Motivation
12
+
13
+ Why is this needed? What problem does it solve?
14
+
15
+ ## Alternatives considered
16
+
17
+ Any other approaches you've thought about?
@@ -0,0 +1,9 @@
1
+ ## Summary
2
+
3
+ What does this PR do?
4
+
5
+ ## Test plan
6
+
7
+ - [ ] Tests pass (`pytest tests/ -v`)
8
+ - [ ] New tests added for new functionality
9
+ - [ ] CHANGELOG.md updated (if user-facing change)
@@ -0,0 +1,20 @@
1
+ name: CI
2
+
3
+ on:
4
+ push:
5
+ branches: [main]
6
+ pull_request:
7
+
8
+ jobs:
9
+ test:
10
+ runs-on: ubuntu-latest
11
+ strategy:
12
+ matrix:
13
+ python-version: ["3.10", "3.11", "3.12"]
14
+ steps:
15
+ - uses: actions/checkout@v4
16
+ - uses: actions/setup-python@v5
17
+ with:
18
+ python-version: ${{ matrix.python-version }}
19
+ - run: pip install -e ".[dev]"
20
+ - run: pytest tests/ -v -x --tb=short -m "not neuron"
@@ -0,0 +1,38 @@
1
+ name: Deploy Docs
2
+
3
+ on:
4
+ push:
5
+ branches: [main]
6
+
7
+ permissions:
8
+ contents: read
9
+ pages: write
10
+ id-token: write
11
+
12
+ concurrency:
13
+ group: pages
14
+ cancel-in-progress: true
15
+
16
+ jobs:
17
+ build:
18
+ runs-on: ubuntu-latest
19
+ steps:
20
+ - uses: actions/checkout@v4
21
+ - uses: actions/setup-python@v5
22
+ with:
23
+ python-version: "3.12"
24
+ - run: pip install mkdocs mkdocs-material
25
+ - run: mkdocs build --strict
26
+ - uses: actions/upload-pages-artifact@v3
27
+ with:
28
+ path: site/
29
+
30
+ deploy:
31
+ needs: build
32
+ runs-on: ubuntu-latest
33
+ environment:
34
+ name: github-pages
35
+ url: ${{ steps.deployment.outputs.page_url }}
36
+ steps:
37
+ - id: deployment
38
+ uses: actions/deploy-pages@v4
@@ -0,0 +1,50 @@
1
+ name: Neuron Hardware Tests
2
+
3
+ on:
4
+ workflow_dispatch:
5
+ inputs:
6
+ instance_type:
7
+ description: "Instance type"
8
+ required: true
9
+ default: "trn1"
10
+ type: choice
11
+ options:
12
+ - trn1
13
+ - trn2
14
+ - inf2
15
+
16
+ env:
17
+ AWS_PROFILE: aws
18
+
19
+ jobs:
20
+ neuron-test:
21
+ runs-on: ubuntu-latest
22
+ permissions:
23
+ id-token: write
24
+ contents: read
25
+ steps:
26
+ - uses: actions/checkout@v4
27
+
28
+ - uses: aws-actions/configure-aws-credentials@v4
29
+ with:
30
+ role-to-assume: ${{ secrets.AWS_ROLE_ARN }}
31
+ aws-region: us-east-1
32
+
33
+ # TODO: Wire to actual Trainium instance via SSM or SSH.
34
+ # The steps below are a scaffold — replace with your instance access method.
35
+ #
36
+ # Example with SSM:
37
+ # - name: Run tests on ${{ inputs.instance_type }}
38
+ # run: |
39
+ # INSTANCE_ID=$(aws ec2 describe-instances \
40
+ # --filters "Name=tag:Name,Values=trnblas-ci-${{ inputs.instance_type }}" \
41
+ # --query 'Reservations[0].Instances[0].InstanceId' --output text)
42
+ # aws ssm send-command --instance-ids $INSTANCE_ID \
43
+ # --document-name "AWS-RunShellScript" \
44
+ # --parameters 'commands=["cd /home/ubuntu/trnblas && git pull && pip install -e \".[neuron,dev]\" && pytest tests/ -v -m neuron --tb=short"]'
45
+ - name: Run neuron tests on ${{ inputs.instance_type }}
46
+ run: |
47
+ echo "Neuron hardware CI scaffold"
48
+ echo "Instance type: ${{ inputs.instance_type }}"
49
+ echo "TODO: Wire to actual ${{ inputs.instance_type }} instance"
50
+ echo "Expected command: pytest tests/ -v -m neuron --tb=short"
@@ -0,0 +1,56 @@
1
+ name: Publish to PyPI
2
+
3
+ on:
4
+ release:
5
+ types: [published]
6
+
7
+ jobs:
8
+ build:
9
+ runs-on: ubuntu-latest
10
+ steps:
11
+ - uses: actions/checkout@v4
12
+ - uses: actions/setup-python@v5
13
+ with:
14
+ python-version: "3.12"
15
+ - run: pip install build
16
+ - run: python -m build
17
+ - uses: actions/upload-artifact@v4
18
+ with:
19
+ name: dist
20
+ path: dist/
21
+
22
+ publish-pypi:
23
+ needs: build
24
+ runs-on: ubuntu-latest
25
+ # Only publish stable releases (v1.2.3), not pre-releases (v1.2.3-rc1)
26
+ if: ${{ !contains(github.event.release.tag_name, '-') }}
27
+ environment:
28
+ name: pypi
29
+ url: https://pypi.org/project/trnblas/
30
+ permissions:
31
+ id-token: write
32
+ steps:
33
+ - uses: actions/download-artifact@v4
34
+ with:
35
+ name: dist
36
+ path: dist/
37
+ - uses: pypa/gh-action-pypi-publish@release/v1
38
+
39
+ publish-testpypi:
40
+ needs: build
41
+ runs-on: ubuntu-latest
42
+ # Only publish pre-releases to TestPyPI
43
+ if: ${{ contains(github.event.release.tag_name, '-') }}
44
+ environment:
45
+ name: testpypi
46
+ url: https://test.pypi.org/project/trnblas/
47
+ permissions:
48
+ id-token: write
49
+ steps:
50
+ - uses: actions/download-artifact@v4
51
+ with:
52
+ name: dist
53
+ path: dist/
54
+ - uses: pypa/gh-action-pypi-publish@release/v1
55
+ with:
56
+ repository-url: https://test.pypi.org/legacy/
@@ -0,0 +1,18 @@
1
+ __pycache__/
2
+ *.py[cod]
3
+ *$py.class
4
+ *.egg-info/
5
+ dist/
6
+ build/
7
+ .eggs/
8
+ *.egg
9
+ .pytest_cache/
10
+ .tox/
11
+ .venv/
12
+ venv/
13
+ *.so
14
+ .idea/
15
+ .vscode/
16
+ *.swp
17
+ *.swo
18
+ *~
@@ -0,0 +1,46 @@
1
+ # Changelog
2
+
3
+ All notable changes to this project will be documented in this file.
4
+
5
+ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
6
+ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
7
+
8
+ ## [Unreleased]
9
+
10
+ ## [0.2.0] - 2026-04-11
11
+
12
+ ### Added
13
+
14
+ - MkDocs Material documentation site at
15
+ [scttfrdmn.github.io/trnblas](https://scttfrdmn.github.io/trnblas/) with
16
+ Installation, Quickstart, API reference (Level 1/2/3, NKI backend), and
17
+ Architecture pages.
18
+ - GitHub Actions CI matrix (Python 3.10, 3.11, 3.12).
19
+ - Neuron hardware CI workflow scaffold (`workflow_dispatch`) — SSM wiring
20
+ deferred until a persistent CI Trainium instance is available.
21
+ - PyPI publishing workflow (OIDC trusted publishers, sdist + wheel on
22
+ release) — matches trnfft pattern.
23
+ - Benchmark suite scaffold (`benchmarks/bench_blas.py`, pytest-benchmark).
24
+ - Issue and PR templates under `.github/`.
25
+ - README badges — CI status, PyPI version, Python versions, License, Docs.
26
+ - Cross-link to trnblas in trnfft's Related Projects table
27
+ (`scttfrdmn/trnfft@7330b3f`).
28
+
29
+ ## [0.1.0] - 2026-04-11
30
+
31
+ ### Added
32
+
33
+ - Level 1 BLAS: `axpy`, `dot`, `nrm2`, `scal`, `asum`, `iamax`.
34
+ - Level 2 BLAS: `gemv`, `symv`, `trmv`, `ger`.
35
+ - Level 3 BLAS: `gemm`, `batched_gemm`, `symm`, `syrk`, `trsm`, `trmm`.
36
+ - NKI dispatch layer with `auto`, `pytorch`, and `nki` backend selection.
37
+ - NKI GEMM kernel stub with stationary tile reuse strategy (scaffolded for
38
+ on-hardware validation on trn1/trn2).
39
+ - DF-MP2 example (`examples/df_mp2.py`) demonstrating the Janesko/TCU use case
40
+ with half-transform GEMMs, Cholesky, triangular solve, and energy evaluation.
41
+ - Test suite covering Level 1/2/3 BLAS correctness against PyTorch/NumPy
42
+ references, with SPD matrix fixtures for symmetric/triangular routines.
43
+
44
+ [Unreleased]: https://github.com/scttfrdmn/trnblas/compare/v0.2.0...HEAD
45
+ [0.2.0]: https://github.com/scttfrdmn/trnblas/releases/tag/v0.2.0
46
+ [0.1.0]: https://github.com/scttfrdmn/trnblas/commits/main
@@ -0,0 +1,105 @@
1
+ # trnblas
2
+
3
+ BLAS operations for AWS Trainium via NKI.
4
+ Part of the trn-* scientific computing suite by Playground Logic.
5
+
6
+ ## What This Is
7
+
8
+ A cuBLAS-equivalent for Trainium. Provides Level 1 (vector), Level 2
9
+ (matrix-vector), and Level 3 (matrix-matrix) BLAS operations with NKI
10
+ kernel acceleration on the Tensor Engine.
11
+
12
+ **Primary use case:** DF-MP2 quantum chemistry with Prof. Ben Janesko at TCU.
13
+ Molecules with >3000 basis functions require sustained GEMM throughput for
14
+ tensor contractions. On 192 Zen4 cores (c8a.48xlarge), a single calculation
15
+ takes ~24 hours at ~$200. trnblas on Trainium targets 3-4x cost reduction
16
+ by exploiting the systolic array for the GEMM-dominated hot path.
17
+
18
+ ## Architecture
19
+
20
+ ```
21
+ trnblas/
22
+ ├── trnblas/
23
+ │ ├── __init__.py # Re-exports all BLAS operations
24
+ │ ├── level1.py # axpy, dot, nrm2, scal, asum, iamax
25
+ │ ├── level2.py # gemv, symv, trmv, ger
26
+ │ ├── level3.py # gemm, batched_gemm, symm, syrk, trsm, trmm
27
+ │ └── nki/
28
+ │ ├── __init__.py
29
+ │ └── dispatch.py # auto/pytorch/nki dispatch + NKI GEMM kernel
30
+ ├── tests/
31
+ │ ├── conftest.py # Fixtures: random matrices, SPD matrices
32
+ │ ├── test_level1.py # Vector operation tests
33
+ │ ├── test_level2.py # Matrix-vector tests
34
+ │ └── test_level3.py # Matrix-matrix tests (GEMM, TRSM, SYRK, etc.)
35
+ ├── examples/
36
+ │ └── df_mp2.py # DF-MP2 energy using trnblas (Janesko use case)
37
+ ├── pyproject.toml
38
+ ├── README.md
39
+ ├── LICENSE # Apache 2.0
40
+ └── CLAUDE.md # This file
41
+ ```
42
+
43
+ ## DF-MP2 → BLAS Mapping
44
+
45
+ The DF-MP2 algorithm maps directly to trnblas Level 3 operations:
46
+
47
+ | DF-MP2 Step | Math | trnblas Call |
48
+ |-------------|------|-------------|
49
+ | Half-transform | (iν\|P) = C_occ^T @ (μν\|P) | `gemm(1.0, C_occ, eri, transA=True)` |
50
+ | MO transform | (ia\|P) = C_vir^T @ (iν\|P) | `gemm(1.0, C_vir, iv_P, transA=True)` |
51
+ | Cholesky | J = L @ L^T | `torch.linalg.cholesky` |
52
+ | Metric solve | L^T @ X = I | `trsm(1.0, L, I, uplo="lower", trans=True)` |
53
+ | Metric contract | B = (ia\|Q) @ J^{-1/2} | `gemm(1.0, ia_P, J_inv_half)` |
54
+ | Energy | T_ab = B_i @ B_j^T | `gemm(1.0, B[i], B[j], transB=True)` |
55
+
56
+ ## NKI GEMM Strategy
57
+
58
+ The NKI GEMM kernel uses stationary tile reuse on the Tensor Engine:
59
+
60
+ 1. Load A tile (128×128) to SBUF as stationary operand
61
+ 2. Stream B tiles through the systolic array as moving operand
62
+ 3. Accumulate partial products in PSUM
63
+ 4. One A load serves all B tiles → 2x fewer SBUF loads vs naive
64
+
65
+ For DF-MP2, the MO coefficient matrix C is the natural stationary operand
66
+ since it's reused across all auxiliary basis indices P.
67
+
68
+ ## Known Gaps & Design Notes
69
+
70
+ - **NKI GEMM is a stub.** Falls back to `torch.matmul` until validated on
71
+ trn1/trn2. The kernel scaffold is in `nki/dispatch.py`.
72
+
73
+ - **No FP64.** Trainium's Tensor Engine maxes out at FP32. For chemistry
74
+ workloads needing higher precision, double-double arithmetic (emulated FP64
75
+ from two FP32 values) is the path — not implemented yet.
76
+
77
+ - **Level 3 dominates.** Level 1/2 are included for API completeness but
78
+ won't get NKI kernels initially. The Tensor Engine is wasted on vector ops.
79
+
80
+ - **DF-MP2 example uses Python loops** over occupied orbital pairs. A
81
+ production implementation would batch these as a single batched_gemm call.
82
+
83
+ ## Dependencies
84
+
85
+ - `torch>=2.1` — tensor operations and CPU fallback
86
+ - `numpy>=1.24` — test reference
87
+ - `neuronxcc` — NKI kernels (optional, only on Neuron hardware)
88
+
89
+ ## Development
90
+
91
+ ```bash
92
+ pip install -e ".[dev]"
93
+ pytest tests/ -v # CPU fallback mode
94
+ python examples/df_mp2.py --demo # Quick DF-MP2 demo
95
+ ```
96
+
97
+ ## Naming Convention
98
+
99
+ Sibling repos in the trn-* suite:
100
+ - `trnfft` — FFT + complex ops (https://github.com/scttfrdmn/trnfft)
101
+ - `trnblas` — BLAS operations (this repo)
102
+ - `trnsolver` — Linear solvers, eigendecomposition (planned)
103
+ - `trnrand` — Random number generation (planned)
104
+
105
+ All repos: Python/NKI, Apache 2.0, Playground Logic.
trnblas-0.2.0/LICENSE ADDED
@@ -0,0 +1,17 @@
1
+ Apache License
2
+ Version 2.0, January 2004
3
+ http://www.apache.org/licenses/
4
+
5
+ Copyright 2026 Playground Logic LLC
6
+
7
+ Licensed under the Apache License, Version 2.0 (the "License");
8
+ you may not use this file except in compliance with the License.
9
+ You may obtain a copy of the License at
10
+
11
+ http://www.apache.org/licenses/LICENSE-2.0
12
+
13
+ Unless required by applicable law or agreed to in writing, software
14
+ distributed under the License is distributed on an "AS IS" BASIS,
15
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
16
+ See the License for the specific language governing permissions and
17
+ limitations under the License.
trnblas-0.2.0/PKG-INFO ADDED
@@ -0,0 +1,139 @@
1
+ Metadata-Version: 2.4
2
+ Name: trnblas
3
+ Version: 0.2.0
4
+ Summary: BLAS operations for AWS Trainium via NKI
5
+ Author-email: Scott Friedman <scott@playgroundlogic.co>
6
+ License-Expression: Apache-2.0
7
+ Project-URL: Homepage, https://github.com/scttfrdmn/trnblas
8
+ Project-URL: Repository, https://github.com/scttfrdmn/trnblas
9
+ Project-URL: Issues, https://github.com/scttfrdmn/trnblas/issues
10
+ Classifier: Development Status :: 3 - Alpha
11
+ Classifier: Intended Audience :: Science/Research
12
+ Classifier: Programming Language :: Python :: 3
13
+ Classifier: Topic :: Scientific/Engineering
14
+ Classifier: Topic :: Scientific/Engineering :: Mathematics
15
+ Requires-Python: >=3.10
16
+ Description-Content-Type: text/markdown
17
+ License-File: LICENSE
18
+ Requires-Dist: torch>=2.1
19
+ Requires-Dist: numpy>=1.24
20
+ Provides-Extra: neuron
21
+ Requires-Dist: neuronxcc>=2.15; extra == "neuron"
22
+ Requires-Dist: torch-neuronx>=2.1; extra == "neuron"
23
+ Provides-Extra: dev
24
+ Requires-Dist: pytest>=7.0; extra == "dev"
25
+ Requires-Dist: pytest-benchmark>=4.0; extra == "dev"
26
+ Requires-Dist: scipy>=1.11; extra == "dev"
27
+ Dynamic: license-file
28
+
29
+ # trnblas
30
+
31
+ [![CI](https://github.com/scttfrdmn/trnblas/actions/workflows/ci.yml/badge.svg)](https://github.com/scttfrdmn/trnblas/actions/workflows/ci.yml)
32
+ [![PyPI](https://img.shields.io/pypi/v/trnblas)](https://pypi.org/project/trnblas/)
33
+ [![Python](https://img.shields.io/pypi/pyversions/trnblas)](https://pypi.org/project/trnblas/)
34
+ [![License](https://img.shields.io/github/license/scttfrdmn/trnblas)](LICENSE)
35
+ [![Docs](https://img.shields.io/badge/docs-mkdocs-blue)](https://scttfrdmn.github.io/trnblas/)
36
+
37
+ BLAS operations for AWS Trainium via NKI (Neuron Kernel Interface).
38
+
39
+ Trainium ships no BLAS library. `trnblas` provides Level 1-3 BLAS operations with NKI kernel acceleration on the Tensor Engine, targeting scientific computing workloads that are GEMM-dominated.
40
+
41
+ Part of the **trn-\*** scientific computing suite by [Playground Logic](https://playgroundlogic.co).
42
+
43
+ ## Why
44
+
45
+ NVIDIA has cuBLAS with 152 optimized routines. Trainium has `torch.matmul`. That's fine for ML training but insufficient for scientific computing codes that need TRSM, SYRK, SYMM, and batched GEMM with specific transpose/scaling semantics.
46
+
47
+ trnblas closes this gap — same BLAS API surface, NKI-accelerated GEMM on Trainium, PyTorch fallback everywhere else.
48
+
49
+ ## Install
50
+
51
+ ```bash
52
+ pip install trnblas
53
+
54
+ # With Neuron hardware support
55
+ pip install trnblas[neuron]
56
+ ```
57
+
58
+ ## Usage
59
+
60
+ ```python
61
+ import torch
62
+ import trnblas
63
+
64
+ # Level 3 — Matrix multiply (the hot path)
65
+ C = trnblas.gemm(alpha=1.0, A=A, B=B, beta=0.5, C=C_init, transA=True)
66
+
67
+ # Batched GEMM (DF-MP2 tensor contractions)
68
+ C = trnblas.batched_gemm(1.0, A_batch, B_batch)
69
+
70
+ # Symmetric matrix multiply (Fock builds)
71
+ F = trnblas.symm(1.0, density, H_core, side="left")
72
+
73
+ # Triangular solve (Cholesky-based density fitting)
74
+ X = trnblas.trsm(1.0, L, B, uplo="lower")
75
+
76
+ # Symmetric rank-k update (metric construction)
77
+ J = trnblas.syrk(1.0, integrals, trans=True)
78
+
79
+ # Level 2 — Matrix-vector
80
+ y = trnblas.gemv(1.0, A, x, beta=1.0, y=y)
81
+
82
+ # Level 1 — Vector operations
83
+ y = trnblas.axpy(alpha, x, y)
84
+ d = trnblas.dot(x, y)
85
+ n = trnblas.nrm2(x)
86
+ ```
87
+
88
+ ## DF-MP2 Example
89
+
90
+ ```bash
91
+ # Run the density-fitted MP2 example (Janesko/TCU use case)
92
+ python examples/df_mp2.py --demo
93
+ python examples/df_mp2.py --nbasis 100 --nocc 20
94
+ ```
95
+
96
+ The example demonstrates all core BLAS operations in a realistic quantum chemistry workflow: Cholesky factorization, triangular solve, half-transform GEMMs, metric contraction, and energy evaluation.
97
+
98
+ ## Operations
99
+
100
+ | Level | Operation | Description |
101
+ |-------|-----------|-------------|
102
+ | 1 | `axpy` | y = αx + y |
103
+ | 1 | `dot` | x^T y |
104
+ | 1 | `nrm2` | ‖x‖₂ |
105
+ | 1 | `scal` | x = αx |
106
+ | 1 | `asum` | Σ\|xᵢ\| |
107
+ | 1 | `iamax` | argmax \|xᵢ\| |
108
+ | 2 | `gemv` | y = α op(A) x + βy |
109
+ | 2 | `symv` | y = α A x + βy (A symmetric) |
110
+ | 2 | `trmv` | x = op(A) x (A triangular) |
111
+ | 2 | `ger` | A = α x yᵀ + A |
112
+ | 3 | `gemm` | C = α op(A) op(B) + βC |
113
+ | 3 | `batched_gemm` | Batched GEMM |
114
+ | 3 | `symm` | C = α A B + βC (A symmetric) |
115
+ | 3 | `syrk` | C = α A Aᵀ + βC |
116
+ | 3 | `trsm` | Solve op(A) X = αB |
117
+ | 3 | `trmm` | B = α op(A) B |
118
+
119
+ ## Status
120
+
121
+ - [x] Level 1-3 BLAS with PyTorch backend
122
+ - [x] GEMM with NKI dispatch stub
123
+ - [x] DF-MP2 example (Janesko/TCU use case)
124
+ - [ ] NKI GEMM kernel validation on trn1/trn2
125
+ - [ ] NKI GEMM with stationary tile reuse
126
+ - [ ] Batched GEMM NKI kernel
127
+ - [ ] Double-double FP64 emulation
128
+ - [ ] Benchmarks vs cuBLAS
129
+
130
+ ## Related Projects
131
+
132
+ | Project | What |
133
+ |---------|------|
134
+ | [trnfft](https://github.com/scttfrdmn/trnfft) | FFT + complex ops for Trainium (Williamson/OSU use case) |
135
+ | trnsolver *(planned)* | Linear solvers and eigendecomposition |
136
+
137
+ ## License
138
+
139
+ Apache 2.0 — Playground Logic LLC
@@ -0,0 +1,111 @@
1
+ # trnblas
2
+
3
+ [![CI](https://github.com/scttfrdmn/trnblas/actions/workflows/ci.yml/badge.svg)](https://github.com/scttfrdmn/trnblas/actions/workflows/ci.yml)
4
+ [![PyPI](https://img.shields.io/pypi/v/trnblas)](https://pypi.org/project/trnblas/)
5
+ [![Python](https://img.shields.io/pypi/pyversions/trnblas)](https://pypi.org/project/trnblas/)
6
+ [![License](https://img.shields.io/github/license/scttfrdmn/trnblas)](LICENSE)
7
+ [![Docs](https://img.shields.io/badge/docs-mkdocs-blue)](https://scttfrdmn.github.io/trnblas/)
8
+
9
+ BLAS operations for AWS Trainium via NKI (Neuron Kernel Interface).
10
+
11
+ Trainium ships no BLAS library. `trnblas` provides Level 1-3 BLAS operations with NKI kernel acceleration on the Tensor Engine, targeting scientific computing workloads that are GEMM-dominated.
12
+
13
+ Part of the **trn-\*** scientific computing suite by [Playground Logic](https://playgroundlogic.co).
14
+
15
+ ## Why
16
+
17
+ NVIDIA has cuBLAS with 152 optimized routines. Trainium has `torch.matmul`. That's fine for ML training but insufficient for scientific computing codes that need TRSM, SYRK, SYMM, and batched GEMM with specific transpose/scaling semantics.
18
+
19
+ trnblas closes this gap — same BLAS API surface, NKI-accelerated GEMM on Trainium, PyTorch fallback everywhere else.
20
+
21
+ ## Install
22
+
23
+ ```bash
24
+ pip install trnblas
25
+
26
+ # With Neuron hardware support
27
+ pip install trnblas[neuron]
28
+ ```
29
+
30
+ ## Usage
31
+
32
+ ```python
33
+ import torch
34
+ import trnblas
35
+
36
+ # Level 3 — Matrix multiply (the hot path)
37
+ C = trnblas.gemm(alpha=1.0, A=A, B=B, beta=0.5, C=C_init, transA=True)
38
+
39
+ # Batched GEMM (DF-MP2 tensor contractions)
40
+ C = trnblas.batched_gemm(1.0, A_batch, B_batch)
41
+
42
+ # Symmetric matrix multiply (Fock builds)
43
+ F = trnblas.symm(1.0, density, H_core, side="left")
44
+
45
+ # Triangular solve (Cholesky-based density fitting)
46
+ X = trnblas.trsm(1.0, L, B, uplo="lower")
47
+
48
+ # Symmetric rank-k update (metric construction)
49
+ J = trnblas.syrk(1.0, integrals, trans=True)
50
+
51
+ # Level 2 — Matrix-vector
52
+ y = trnblas.gemv(1.0, A, x, beta=1.0, y=y)
53
+
54
+ # Level 1 — Vector operations
55
+ y = trnblas.axpy(alpha, x, y)
56
+ d = trnblas.dot(x, y)
57
+ n = trnblas.nrm2(x)
58
+ ```
59
+
60
+ ## DF-MP2 Example
61
+
62
+ ```bash
63
+ # Run the density-fitted MP2 example (Janesko/TCU use case)
64
+ python examples/df_mp2.py --demo
65
+ python examples/df_mp2.py --nbasis 100 --nocc 20
66
+ ```
67
+
68
+ The example demonstrates all core BLAS operations in a realistic quantum chemistry workflow: Cholesky factorization, triangular solve, half-transform GEMMs, metric contraction, and energy evaluation.
69
+
70
+ ## Operations
71
+
72
+ | Level | Operation | Description |
73
+ |-------|-----------|-------------|
74
+ | 1 | `axpy` | y = αx + y |
75
+ | 1 | `dot` | x^T y |
76
+ | 1 | `nrm2` | ‖x‖₂ |
77
+ | 1 | `scal` | x = αx |
78
+ | 1 | `asum` | Σ\|xᵢ\| |
79
+ | 1 | `iamax` | argmax \|xᵢ\| |
80
+ | 2 | `gemv` | y = α op(A) x + βy |
81
+ | 2 | `symv` | y = α A x + βy (A symmetric) |
82
+ | 2 | `trmv` | x = op(A) x (A triangular) |
83
+ | 2 | `ger` | A = α x yᵀ + A |
84
+ | 3 | `gemm` | C = α op(A) op(B) + βC |
85
+ | 3 | `batched_gemm` | Batched GEMM |
86
+ | 3 | `symm` | C = α A B + βC (A symmetric) |
87
+ | 3 | `syrk` | C = α A Aᵀ + βC |
88
+ | 3 | `trsm` | Solve op(A) X = αB |
89
+ | 3 | `trmm` | B = α op(A) B |
90
+
91
+ ## Status
92
+
93
+ - [x] Level 1-3 BLAS with PyTorch backend
94
+ - [x] GEMM with NKI dispatch stub
95
+ - [x] DF-MP2 example (Janesko/TCU use case)
96
+ - [ ] NKI GEMM kernel validation on trn1/trn2
97
+ - [ ] NKI GEMM with stationary tile reuse
98
+ - [ ] Batched GEMM NKI kernel
99
+ - [ ] Double-double FP64 emulation
100
+ - [ ] Benchmarks vs cuBLAS
101
+
102
+ ## Related Projects
103
+
104
+ | Project | What |
105
+ |---------|------|
106
+ | [trnfft](https://github.com/scttfrdmn/trnfft) | FFT + complex ops for Trainium (Williamson/OSU use case) |
107
+ | trnsolver *(planned)* | Linear solvers and eigendecomposition |
108
+
109
+ ## License
110
+
111
+ Apache 2.0 — Playground Logic LLC