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.
- trnblas-0.2.0/.github/ISSUE_TEMPLATE/bug_report.md +31 -0
- trnblas-0.2.0/.github/ISSUE_TEMPLATE/feature_request.md +17 -0
- trnblas-0.2.0/.github/pull_request_template.md +9 -0
- trnblas-0.2.0/.github/workflows/ci.yml +20 -0
- trnblas-0.2.0/.github/workflows/docs.yml +38 -0
- trnblas-0.2.0/.github/workflows/neuron.yml +50 -0
- trnblas-0.2.0/.github/workflows/publish.yml +56 -0
- trnblas-0.2.0/.gitignore +18 -0
- trnblas-0.2.0/CHANGELOG.md +46 -0
- trnblas-0.2.0/CLAUDE.md +105 -0
- trnblas-0.2.0/LICENSE +17 -0
- trnblas-0.2.0/PKG-INFO +139 -0
- trnblas-0.2.0/README.md +111 -0
- trnblas-0.2.0/benchmarks/bench_blas.py +38 -0
- trnblas-0.2.0/benchmarks/conftest.py +16 -0
- trnblas-0.2.0/docs/api/level1.md +28 -0
- trnblas-0.2.0/docs/api/level2.md +20 -0
- trnblas-0.2.0/docs/api/level3.md +32 -0
- trnblas-0.2.0/docs/api/nki.md +29 -0
- trnblas-0.2.0/docs/architecture.md +52 -0
- trnblas-0.2.0/docs/index.md +31 -0
- trnblas-0.2.0/docs/installation.md +33 -0
- trnblas-0.2.0/docs/quickstart.md +53 -0
- trnblas-0.2.0/examples/df_mp2.py +147 -0
- trnblas-0.2.0/mkdocs.yml +32 -0
- trnblas-0.2.0/pyproject.toml +50 -0
- trnblas-0.2.0/setup.cfg +4 -0
- trnblas-0.2.0/tests/conftest.py +42 -0
- trnblas-0.2.0/tests/test_level1.py +62 -0
- trnblas-0.2.0/tests/test_level2.py +78 -0
- trnblas-0.2.0/tests/test_level3.py +151 -0
- trnblas-0.2.0/trnblas/__init__.py +35 -0
- trnblas-0.2.0/trnblas/level1.py +45 -0
- trnblas-0.2.0/trnblas/level2.py +86 -0
- trnblas-0.2.0/trnblas/level3.py +203 -0
- trnblas-0.2.0/trnblas/nki/__init__.py +5 -0
- trnblas-0.2.0/trnblas/nki/dispatch.py +121 -0
- trnblas-0.2.0/trnblas.egg-info/PKG-INFO +139 -0
- trnblas-0.2.0/trnblas.egg-info/SOURCES.txt +40 -0
- trnblas-0.2.0/trnblas.egg-info/dependency_links.txt +1 -0
- trnblas-0.2.0/trnblas.egg-info/requires.txt +11 -0
- 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,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/
|
trnblas-0.2.0/.gitignore
ADDED
|
@@ -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
|
trnblas-0.2.0/CLAUDE.md
ADDED
|
@@ -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
|
+
[](https://github.com/scttfrdmn/trnblas/actions/workflows/ci.yml)
|
|
32
|
+
[](https://pypi.org/project/trnblas/)
|
|
33
|
+
[](https://pypi.org/project/trnblas/)
|
|
34
|
+
[](LICENSE)
|
|
35
|
+
[](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
|
trnblas-0.2.0/README.md
ADDED
|
@@ -0,0 +1,111 @@
|
|
|
1
|
+
# trnblas
|
|
2
|
+
|
|
3
|
+
[](https://github.com/scttfrdmn/trnblas/actions/workflows/ci.yml)
|
|
4
|
+
[](https://pypi.org/project/trnblas/)
|
|
5
|
+
[](https://pypi.org/project/trnblas/)
|
|
6
|
+
[](LICENSE)
|
|
7
|
+
[](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
|