swage-compiler 0.5.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.
- swage_compiler-0.5.0/.gitignore +31 -0
- swage_compiler-0.5.0/LICENSE +21 -0
- swage_compiler-0.5.0/PKG-INFO +227 -0
- swage_compiler-0.5.0/README.md +199 -0
- swage_compiler-0.5.0/pyproject.toml +59 -0
- swage_compiler-0.5.0/python/swage/__init__.py +8 -0
- swage_compiler-0.5.0/python/swage/_frontend.py +755 -0
- swage_compiler-0.5.0/python/swage/_runtime.py +561 -0
- swage_compiler-0.5.0/python/swage/_segmented_qualification.py +774 -0
- swage_compiler-0.5.0/python/swage/env.py +88 -0
- swage_compiler-0.5.0/python/swage/language.py +86 -0
|
@@ -0,0 +1,31 @@
|
|
|
1
|
+
# Build trees
|
|
2
|
+
build/
|
|
3
|
+
build-*/
|
|
4
|
+
*.o
|
|
5
|
+
*.a
|
|
6
|
+
*.so
|
|
7
|
+
|
|
8
|
+
# Python
|
|
9
|
+
__pycache__/
|
|
10
|
+
*.py[cod]
|
|
11
|
+
*.egg-info/
|
|
12
|
+
dist/
|
|
13
|
+
.venv/
|
|
14
|
+
venv/
|
|
15
|
+
|
|
16
|
+
# Tool caches
|
|
17
|
+
.cache/
|
|
18
|
+
.pytest_cache/
|
|
19
|
+
.ruff_cache/
|
|
20
|
+
.mypy_cache/
|
|
21
|
+
|
|
22
|
+
# Editors
|
|
23
|
+
.vscode/
|
|
24
|
+
.idea/
|
|
25
|
+
*.swp
|
|
26
|
+
|
|
27
|
+
# Docs output
|
|
28
|
+
site/
|
|
29
|
+
|
|
30
|
+
# Swage artifacts
|
|
31
|
+
.swage-cache/
|
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2026 Abhik Sarkar and Swage contributors
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
|
@@ -0,0 +1,227 @@
|
|
|
1
|
+
Metadata-Version: 2.5
|
|
2
|
+
Name: swage-compiler
|
|
3
|
+
Version: 0.5.0
|
|
4
|
+
Summary: Python-embedded MLIR/LLVM GPU compiler for variable-sized dense segments
|
|
5
|
+
Project-URL: Repository, https://github.com/abhiksark/swage
|
|
6
|
+
Author-email: Abhik Sarkar <abhiksark@gmail.com>
|
|
7
|
+
License-Expression: MIT
|
|
8
|
+
License-File: LICENSE
|
|
9
|
+
Keywords: compiler,cuda,gpu,llvm,mlir,ragged,segments
|
|
10
|
+
Classifier: Development Status :: 2 - Pre-Alpha
|
|
11
|
+
Classifier: Intended Audience :: Developers
|
|
12
|
+
Classifier: Intended Audience :: Science/Research
|
|
13
|
+
Classifier: Programming Language :: Python :: 3
|
|
14
|
+
Classifier: Topic :: Software Development :: Compilers
|
|
15
|
+
Requires-Python: >=3.10
|
|
16
|
+
Provides-Extra: dev
|
|
17
|
+
Requires-Dist: build; extra == 'dev'
|
|
18
|
+
Requires-Dist: hatchling; extra == 'dev'
|
|
19
|
+
Requires-Dist: lit; extra == 'dev'
|
|
20
|
+
Requires-Dist: pytest; extra == 'dev'
|
|
21
|
+
Requires-Dist: ruff; extra == 'dev'
|
|
22
|
+
Provides-Extra: docs
|
|
23
|
+
Requires-Dist: mkdocs; extra == 'docs'
|
|
24
|
+
Requires-Dist: mkdocs-material; extra == 'docs'
|
|
25
|
+
Provides-Extra: pytorch
|
|
26
|
+
Requires-Dist: torch>=2.6; extra == 'pytorch'
|
|
27
|
+
Description-Content-Type: text/markdown
|
|
28
|
+
|
|
29
|
+
# Swage
|
|
30
|
+
|
|
31
|
+
[](https://github.com/abhiksark/swage/actions/workflows/ci-python.yml)
|
|
32
|
+
[](https://github.com/abhiksark/swage/actions/workflows/ci-cpp.yml)
|
|
33
|
+
[](https://github.com/abhiksark/swage/actions/workflows/ci-gpu.yml)
|
|
34
|
+
|
|
35
|
+
**Turn variable-sized dense segments into efficient GPU tile tasks.**
|
|
36
|
+
|
|
37
|
+
Swage is an experimental Python-embedded GPU compiler built on MLIR and
|
|
38
|
+
LLVM. You write a Triton-like kernel that describes what happens to *one
|
|
39
|
+
logical segment*; Swage is designed to derive fixed-size GPU tile tasks and
|
|
40
|
+
generate NVIDIA GPU code through MLIR, LLVM, and NVPTX.
|
|
41
|
+
|
|
42
|
+
> **Status: pre-alpha, M7 internal mixed-policy qualification complete.** The
|
|
43
|
+
> canonical fixed vector-add kernel remains the public execution subset. Internal
|
|
44
|
+
> segmented sum, max, and stable ragged-softmax modules lower to a sequential
|
|
45
|
+
> CPU oracle and one CTA per segment on NVIDIA GPUs. For one canonical identity
|
|
46
|
+
> segmented sum, a private M7 path materializes warp and CTA task IDs and runs
|
|
47
|
+
> pure or fused mixed schedules. Public segment syntax and public segmented
|
|
48
|
+
> launch remain planned. See [Current status](#current-status) for the exact
|
|
49
|
+
> boundary.
|
|
50
|
+
|
|
51
|
+
## The programming model
|
|
52
|
+
|
|
53
|
+
The executable fixed-block vector-add subset is deliberately Triton-like. It
|
|
54
|
+
infers its narrow ABI from PyTorch CUDA tensors and Python integers (`n` maps
|
|
55
|
+
to `sl.int32`):
|
|
56
|
+
|
|
57
|
+
```python
|
|
58
|
+
import swage as sw
|
|
59
|
+
import swage.language as sl
|
|
60
|
+
import torch
|
|
61
|
+
|
|
62
|
+
|
|
63
|
+
@sw.jit
|
|
64
|
+
def add_kernel(x_ptr, y_ptr, output_ptr, n, BLOCK: sl.constexpr):
|
|
65
|
+
pid = sl.program_id(0)
|
|
66
|
+
offsets = pid * BLOCK + sl.arange(0, BLOCK)
|
|
67
|
+
mask = offsets < n
|
|
68
|
+
x = sl.load(x_ptr + offsets, mask=mask, other=0.0)
|
|
69
|
+
y = sl.load(y_ptr + offsets, mask=mask, other=0.0)
|
|
70
|
+
sl.store(output_ptr + offsets, x + y, mask=mask)
|
|
71
|
+
|
|
72
|
+
|
|
73
|
+
length = 1024
|
|
74
|
+
block = 128
|
|
75
|
+
x = torch.randn(length, device="cuda", dtype=torch.float32)
|
|
76
|
+
y = torch.randn(length, device="cuda", dtype=torch.float32)
|
|
77
|
+
output = torch.empty_like(x)
|
|
78
|
+
add_kernel.launch(
|
|
79
|
+
arguments={
|
|
80
|
+
"x_ptr": x,
|
|
81
|
+
"y_ptr": y,
|
|
82
|
+
"output_ptr": output,
|
|
83
|
+
"n": length,
|
|
84
|
+
},
|
|
85
|
+
constexprs={"BLOCK": block},
|
|
86
|
+
grid=((length + block - 1) // block,),
|
|
87
|
+
)
|
|
88
|
+
torch.testing.assert_close(output, x + y)
|
|
89
|
+
```
|
|
90
|
+
|
|
91
|
+
The PyPI wheel installs the importable `swage` package. It does not contain
|
|
92
|
+
the native `mlir_swage` package or compiler build output. Both `launch()` and
|
|
93
|
+
`emit_mlir()` require a bindings-enabled native build and its build-tree
|
|
94
|
+
`mlir_swage` package. Install the optional PyTorch dependency with
|
|
95
|
+
`pip install "swage-compiler[pytorch]"`. `launch()` is keyword-only,
|
|
96
|
+
asynchronous, returns `None`, and uses the current PyTorch CUDA stream.
|
|
97
|
+
`emit_mlir()` remains compile-only, and its explicit `signature=` form
|
|
98
|
+
remains available without PyTorch. Direct kernel calls remain unavailable.
|
|
99
|
+
PTX emission is an internal runtime operation, not a public API.
|
|
100
|
+
|
|
101
|
+
The research target is the segment API: one segment-local program, from
|
|
102
|
+
which the compiler derives packing, bucketing, partitioning, partial
|
|
103
|
+
reductions, and static or persistent scheduling as the runtime
|
|
104
|
+
segment-length distribution changes:
|
|
105
|
+
|
|
106
|
+
```python
|
|
107
|
+
@sw.jit
|
|
108
|
+
def segmented_softmax(values_ptr, offsets_ptr, output_ptr):
|
|
109
|
+
sid = sl.segment_id(0)
|
|
110
|
+
segment = sl.segment(values_ptr, offsets_ptr, sid)
|
|
111
|
+
|
|
112
|
+
x = sl.load_segment(segment)
|
|
113
|
+
maximum = sl.max(x)
|
|
114
|
+
numerator = sl.exp(x - maximum)
|
|
115
|
+
denominator = sl.sum(numerator)
|
|
116
|
+
|
|
117
|
+
sl.store_segment(output_ptr, segment, numerator / denominator)
|
|
118
|
+
```
|
|
119
|
+
|
|
120
|
+
## Architecture
|
|
121
|
+
|
|
122
|
+
```text
|
|
123
|
+
Python @sw.jit kernel
|
|
124
|
+
│ restricted Python AST
|
|
125
|
+
▼
|
|
126
|
+
Swage semantic MLIR
|
|
127
|
+
├── M3 fixed vector add → gpu / nvvm / LLVM IR
|
|
128
|
+
│ → PTX → CUDA Driver API
|
|
129
|
+
│ → current PyTorch stream
|
|
130
|
+
├── M4 segmented sum/max → sequential scf/memref CPU oracle
|
|
131
|
+
│ → one CTA per segment → NVPTX
|
|
132
|
+
├── M5 ragged softmax → max → exponential sum → normalize/store
|
|
133
|
+
│ → internal CPU and one-CTA GPU qualification
|
|
134
|
+
├── M6 identity segmented sum → SwagePlan classify
|
|
135
|
+
└── M7 private materialization → stable warp/CTA task IDs
|
|
136
|
+
→ pure warp or CTA task kernels
|
|
137
|
+
→ one fused mixed GPU kernel
|
|
138
|
+
```
|
|
139
|
+
|
|
140
|
+
The three-level vocabulary is load-bearing: a **segment** is a logical,
|
|
141
|
+
runtime-sized, internally dense data object; a **task** is a schedulable
|
|
142
|
+
unit of execution; a **tile** is a fixed-size physical unit processed by a
|
|
143
|
+
warp or CTA. See
|
|
144
|
+
[docs/concepts/segments-tiles-tasks.md](docs/concepts/segments-tiles-tasks.md).
|
|
145
|
+
|
|
146
|
+
## Current status
|
|
147
|
+
|
|
148
|
+
| Stage | State |
|
|
149
|
+
|---|---|
|
|
150
|
+
| `swage` MLIR dialect (`!swage.segment<T>`, `segment_id`, `make_segment`, `extent`) | **Works today**: parses, prints, verifies; lit-tested |
|
|
151
|
+
| Pinned out-of-tree LLVM/MLIR build (`llvmorg-22.1.8`) + `swage-opt` | **Works today** |
|
|
152
|
+
| `python -m swage.env` environment diagnostics | **Works today** |
|
|
153
|
+
| Native `mlir_swage` bindings package | **Works today** from the build tree; integration-tested |
|
|
154
|
+
| Python AST → verified live `mlir_swage.ir.Module` (fixed-block vector add, inferred or explicit signature) | **Works today**; `emit_mlir()` remains compile-only |
|
|
155
|
+
| Fixed vector add lowering through LLVM NVPTX to deterministic PTX | **Works today, internal**; native-tested for exact targets |
|
|
156
|
+
| CUDA Driver launch, cache, and real GPU result | **Works today for the public M3 subset**; trusted A6000 GPU workflow |
|
|
157
|
+
| Native segmented sum/max lowering | **Works today for canonical internal qualification modules**; upstream `mlir-runner` CPU oracle and one-CTA `sm_86` tests |
|
|
158
|
+
| Native stable ragged-softmax lowering | **Works today for the canonical internal qualification module**; max, exponential-sum, and normalization/store phases match PyTorch and the CPU oracle on RTX A6000 `sm_86` |
|
|
159
|
+
| `swage_plan` policy attribute, task-range type, classify operation, and `--swage-to-plan` | **Works today, internal** for one canonical identity segmented sum; the semantic function is preserved and unsupported inputs fail before mutation |
|
|
160
|
+
| Internal host task descriptor generation | **Works today, unit-tested** for validated i32 metadata; emits one stable warp or CTA descriptor per segment |
|
|
161
|
+
| Internal M7 planned segmented-sum execution | **Works today for canonical qualification only**; a private path consumes the planning threshold, materializes stable task IDs, and executes pure 32-thread warp, pure 128-thread CTA, or one-launch fused mixed schedules on NVIDIA GPUs |
|
|
162
|
+
| Public segment frontend and segmented launch, including ragged softmax | Planned |
|
|
163
|
+
| Packed warps, split CTAs, partial reductions, queues, persistent scheduling, and broader policies | Planned for M8 and later milestones |
|
|
164
|
+
|
|
165
|
+
The M5 differential suite covers all-empty, all-singleton, many-tiny,
|
|
166
|
+
few-huge, one-outlier, and alternating-empty segment distributions. The
|
|
167
|
+
runner is internal: it retains the five-argument values, offsets, output,
|
|
168
|
+
value-count, and segment-count ABI, validates host-visible metadata, and does
|
|
169
|
+
not widen `swage.language`, `emit_mlir()`, or public `launch()`.
|
|
170
|
+
|
|
171
|
+
The M6 planning conversion admits only the capture-free, map-free,
|
|
172
|
+
single-stage identity segmented sum. It adds a private planning companion
|
|
173
|
+
that records warp then CTA as the legal policy order. M7 privately clones the
|
|
174
|
+
semantic module, runs that conversion, reads the recorded threshold, and
|
|
175
|
+
classifies validated runtime metadata into stable warp and CTA segment-ID
|
|
176
|
+
lists. Pure policies use the same task-ID ABI. Mixed execution uses one
|
|
177
|
+
128-thread kernel with four one-segment warp slots per warp block followed by
|
|
178
|
+
one block per CTA task.
|
|
179
|
+
|
|
180
|
+
The frozen `bimodal` benchmark on NVIDIA RTX A6000 `sm_86` measured medians of
|
|
181
|
+
`0.067584 ms` for pure warp, `0.070656 ms` for pure CTA, and `0.063488 ms` for
|
|
182
|
+
mixed execution. Its mixed-to-best-pure ratio is `0.939394`, passing the
|
|
183
|
+
predeclared maximum of `1.05`. This qualifies only the private canonical
|
|
184
|
+
identity-sum path. It does not add public segment syntax or public segmented
|
|
185
|
+
execution.
|
|
186
|
+
|
|
187
|
+
The research question: *can one segment-local program automatically produce
|
|
188
|
+
competitive warp, CTA, split-CTA, and persistent schedules as the runtime
|
|
189
|
+
segment-length distribution changes?* Swage did not invent ragged tensors,
|
|
190
|
+
segmented reductions, persistent kernels, or tile programming; the intended
|
|
191
|
+
contribution is the automatic derivation of the schedule from one
|
|
192
|
+
segment-local kernel.
|
|
193
|
+
|
|
194
|
+
## Getting started
|
|
195
|
+
|
|
196
|
+
```bash
|
|
197
|
+
# CPU-only: build the pinned LLVM/MLIR (once, ~1 hour), then Swage
|
|
198
|
+
./scripts/fetch_llvm.sh
|
|
199
|
+
./scripts/build_llvm.sh
|
|
200
|
+
./scripts/build_swage.sh # builds swage-opt and runs the lit suite
|
|
201
|
+
|
|
202
|
+
# Python package and tests
|
|
203
|
+
make setup
|
|
204
|
+
make test
|
|
205
|
+
```
|
|
206
|
+
|
|
207
|
+
The native bindings require the pinned LLVM/MLIR install with its Python
|
|
208
|
+
bindings enabled. Run `ninja -C build check-swage-python` after the native
|
|
209
|
+
build; it sets the build-tree `PYTHONPATH` for `mlir_swage`. On Linux with
|
|
210
|
+
PyTorch CUDA and `libcuda`, the M3 subset can then compile and launch the
|
|
211
|
+
fixed vector add. The internal native suite also qualifies M4 segmented sum
|
|
212
|
+
and max plus M5 ragged softmax. The public launch path has no CUDA toolkit
|
|
213
|
+
dependency.
|
|
214
|
+
|
|
215
|
+
See the [runnable M3 walkthrough](docs/quickstart.md#execute-fixed-vector-add).
|
|
216
|
+
A GPU is not required for the CPU and compile-only development paths.
|
|
217
|
+
|
|
218
|
+
## More
|
|
219
|
+
|
|
220
|
+
- [DESIGN.md](DESIGN.md): architecture and design invariants
|
|
221
|
+
- [ROADMAP.md](ROADMAP.md): phased plan and honest phase status
|
|
222
|
+
- [CONTRIBUTING.md](CONTRIBUTING.md): contributor paths, CPU-only onboarding
|
|
223
|
+
- [docs/adr/](docs/adr/): architecture decision records
|
|
224
|
+
|
|
225
|
+
## License
|
|
226
|
+
|
|
227
|
+
MIT. See [LICENSE](LICENSE).
|
|
@@ -0,0 +1,199 @@
|
|
|
1
|
+
# Swage
|
|
2
|
+
|
|
3
|
+
[](https://github.com/abhiksark/swage/actions/workflows/ci-python.yml)
|
|
4
|
+
[](https://github.com/abhiksark/swage/actions/workflows/ci-cpp.yml)
|
|
5
|
+
[](https://github.com/abhiksark/swage/actions/workflows/ci-gpu.yml)
|
|
6
|
+
|
|
7
|
+
**Turn variable-sized dense segments into efficient GPU tile tasks.**
|
|
8
|
+
|
|
9
|
+
Swage is an experimental Python-embedded GPU compiler built on MLIR and
|
|
10
|
+
LLVM. You write a Triton-like kernel that describes what happens to *one
|
|
11
|
+
logical segment*; Swage is designed to derive fixed-size GPU tile tasks and
|
|
12
|
+
generate NVIDIA GPU code through MLIR, LLVM, and NVPTX.
|
|
13
|
+
|
|
14
|
+
> **Status: pre-alpha, M7 internal mixed-policy qualification complete.** The
|
|
15
|
+
> canonical fixed vector-add kernel remains the public execution subset. Internal
|
|
16
|
+
> segmented sum, max, and stable ragged-softmax modules lower to a sequential
|
|
17
|
+
> CPU oracle and one CTA per segment on NVIDIA GPUs. For one canonical identity
|
|
18
|
+
> segmented sum, a private M7 path materializes warp and CTA task IDs and runs
|
|
19
|
+
> pure or fused mixed schedules. Public segment syntax and public segmented
|
|
20
|
+
> launch remain planned. See [Current status](#current-status) for the exact
|
|
21
|
+
> boundary.
|
|
22
|
+
|
|
23
|
+
## The programming model
|
|
24
|
+
|
|
25
|
+
The executable fixed-block vector-add subset is deliberately Triton-like. It
|
|
26
|
+
infers its narrow ABI from PyTorch CUDA tensors and Python integers (`n` maps
|
|
27
|
+
to `sl.int32`):
|
|
28
|
+
|
|
29
|
+
```python
|
|
30
|
+
import swage as sw
|
|
31
|
+
import swage.language as sl
|
|
32
|
+
import torch
|
|
33
|
+
|
|
34
|
+
|
|
35
|
+
@sw.jit
|
|
36
|
+
def add_kernel(x_ptr, y_ptr, output_ptr, n, BLOCK: sl.constexpr):
|
|
37
|
+
pid = sl.program_id(0)
|
|
38
|
+
offsets = pid * BLOCK + sl.arange(0, BLOCK)
|
|
39
|
+
mask = offsets < n
|
|
40
|
+
x = sl.load(x_ptr + offsets, mask=mask, other=0.0)
|
|
41
|
+
y = sl.load(y_ptr + offsets, mask=mask, other=0.0)
|
|
42
|
+
sl.store(output_ptr + offsets, x + y, mask=mask)
|
|
43
|
+
|
|
44
|
+
|
|
45
|
+
length = 1024
|
|
46
|
+
block = 128
|
|
47
|
+
x = torch.randn(length, device="cuda", dtype=torch.float32)
|
|
48
|
+
y = torch.randn(length, device="cuda", dtype=torch.float32)
|
|
49
|
+
output = torch.empty_like(x)
|
|
50
|
+
add_kernel.launch(
|
|
51
|
+
arguments={
|
|
52
|
+
"x_ptr": x,
|
|
53
|
+
"y_ptr": y,
|
|
54
|
+
"output_ptr": output,
|
|
55
|
+
"n": length,
|
|
56
|
+
},
|
|
57
|
+
constexprs={"BLOCK": block},
|
|
58
|
+
grid=((length + block - 1) // block,),
|
|
59
|
+
)
|
|
60
|
+
torch.testing.assert_close(output, x + y)
|
|
61
|
+
```
|
|
62
|
+
|
|
63
|
+
The PyPI wheel installs the importable `swage` package. It does not contain
|
|
64
|
+
the native `mlir_swage` package or compiler build output. Both `launch()` and
|
|
65
|
+
`emit_mlir()` require a bindings-enabled native build and its build-tree
|
|
66
|
+
`mlir_swage` package. Install the optional PyTorch dependency with
|
|
67
|
+
`pip install "swage-compiler[pytorch]"`. `launch()` is keyword-only,
|
|
68
|
+
asynchronous, returns `None`, and uses the current PyTorch CUDA stream.
|
|
69
|
+
`emit_mlir()` remains compile-only, and its explicit `signature=` form
|
|
70
|
+
remains available without PyTorch. Direct kernel calls remain unavailable.
|
|
71
|
+
PTX emission is an internal runtime operation, not a public API.
|
|
72
|
+
|
|
73
|
+
The research target is the segment API: one segment-local program, from
|
|
74
|
+
which the compiler derives packing, bucketing, partitioning, partial
|
|
75
|
+
reductions, and static or persistent scheduling as the runtime
|
|
76
|
+
segment-length distribution changes:
|
|
77
|
+
|
|
78
|
+
```python
|
|
79
|
+
@sw.jit
|
|
80
|
+
def segmented_softmax(values_ptr, offsets_ptr, output_ptr):
|
|
81
|
+
sid = sl.segment_id(0)
|
|
82
|
+
segment = sl.segment(values_ptr, offsets_ptr, sid)
|
|
83
|
+
|
|
84
|
+
x = sl.load_segment(segment)
|
|
85
|
+
maximum = sl.max(x)
|
|
86
|
+
numerator = sl.exp(x - maximum)
|
|
87
|
+
denominator = sl.sum(numerator)
|
|
88
|
+
|
|
89
|
+
sl.store_segment(output_ptr, segment, numerator / denominator)
|
|
90
|
+
```
|
|
91
|
+
|
|
92
|
+
## Architecture
|
|
93
|
+
|
|
94
|
+
```text
|
|
95
|
+
Python @sw.jit kernel
|
|
96
|
+
│ restricted Python AST
|
|
97
|
+
▼
|
|
98
|
+
Swage semantic MLIR
|
|
99
|
+
├── M3 fixed vector add → gpu / nvvm / LLVM IR
|
|
100
|
+
│ → PTX → CUDA Driver API
|
|
101
|
+
│ → current PyTorch stream
|
|
102
|
+
├── M4 segmented sum/max → sequential scf/memref CPU oracle
|
|
103
|
+
│ → one CTA per segment → NVPTX
|
|
104
|
+
├── M5 ragged softmax → max → exponential sum → normalize/store
|
|
105
|
+
│ → internal CPU and one-CTA GPU qualification
|
|
106
|
+
├── M6 identity segmented sum → SwagePlan classify
|
|
107
|
+
└── M7 private materialization → stable warp/CTA task IDs
|
|
108
|
+
→ pure warp or CTA task kernels
|
|
109
|
+
→ one fused mixed GPU kernel
|
|
110
|
+
```
|
|
111
|
+
|
|
112
|
+
The three-level vocabulary is load-bearing: a **segment** is a logical,
|
|
113
|
+
runtime-sized, internally dense data object; a **task** is a schedulable
|
|
114
|
+
unit of execution; a **tile** is a fixed-size physical unit processed by a
|
|
115
|
+
warp or CTA. See
|
|
116
|
+
[docs/concepts/segments-tiles-tasks.md](docs/concepts/segments-tiles-tasks.md).
|
|
117
|
+
|
|
118
|
+
## Current status
|
|
119
|
+
|
|
120
|
+
| Stage | State |
|
|
121
|
+
|---|---|
|
|
122
|
+
| `swage` MLIR dialect (`!swage.segment<T>`, `segment_id`, `make_segment`, `extent`) | **Works today**: parses, prints, verifies; lit-tested |
|
|
123
|
+
| Pinned out-of-tree LLVM/MLIR build (`llvmorg-22.1.8`) + `swage-opt` | **Works today** |
|
|
124
|
+
| `python -m swage.env` environment diagnostics | **Works today** |
|
|
125
|
+
| Native `mlir_swage` bindings package | **Works today** from the build tree; integration-tested |
|
|
126
|
+
| Python AST → verified live `mlir_swage.ir.Module` (fixed-block vector add, inferred or explicit signature) | **Works today**; `emit_mlir()` remains compile-only |
|
|
127
|
+
| Fixed vector add lowering through LLVM NVPTX to deterministic PTX | **Works today, internal**; native-tested for exact targets |
|
|
128
|
+
| CUDA Driver launch, cache, and real GPU result | **Works today for the public M3 subset**; trusted A6000 GPU workflow |
|
|
129
|
+
| Native segmented sum/max lowering | **Works today for canonical internal qualification modules**; upstream `mlir-runner` CPU oracle and one-CTA `sm_86` tests |
|
|
130
|
+
| Native stable ragged-softmax lowering | **Works today for the canonical internal qualification module**; max, exponential-sum, and normalization/store phases match PyTorch and the CPU oracle on RTX A6000 `sm_86` |
|
|
131
|
+
| `swage_plan` policy attribute, task-range type, classify operation, and `--swage-to-plan` | **Works today, internal** for one canonical identity segmented sum; the semantic function is preserved and unsupported inputs fail before mutation |
|
|
132
|
+
| Internal host task descriptor generation | **Works today, unit-tested** for validated i32 metadata; emits one stable warp or CTA descriptor per segment |
|
|
133
|
+
| Internal M7 planned segmented-sum execution | **Works today for canonical qualification only**; a private path consumes the planning threshold, materializes stable task IDs, and executes pure 32-thread warp, pure 128-thread CTA, or one-launch fused mixed schedules on NVIDIA GPUs |
|
|
134
|
+
| Public segment frontend and segmented launch, including ragged softmax | Planned |
|
|
135
|
+
| Packed warps, split CTAs, partial reductions, queues, persistent scheduling, and broader policies | Planned for M8 and later milestones |
|
|
136
|
+
|
|
137
|
+
The M5 differential suite covers all-empty, all-singleton, many-tiny,
|
|
138
|
+
few-huge, one-outlier, and alternating-empty segment distributions. The
|
|
139
|
+
runner is internal: it retains the five-argument values, offsets, output,
|
|
140
|
+
value-count, and segment-count ABI, validates host-visible metadata, and does
|
|
141
|
+
not widen `swage.language`, `emit_mlir()`, or public `launch()`.
|
|
142
|
+
|
|
143
|
+
The M6 planning conversion admits only the capture-free, map-free,
|
|
144
|
+
single-stage identity segmented sum. It adds a private planning companion
|
|
145
|
+
that records warp then CTA as the legal policy order. M7 privately clones the
|
|
146
|
+
semantic module, runs that conversion, reads the recorded threshold, and
|
|
147
|
+
classifies validated runtime metadata into stable warp and CTA segment-ID
|
|
148
|
+
lists. Pure policies use the same task-ID ABI. Mixed execution uses one
|
|
149
|
+
128-thread kernel with four one-segment warp slots per warp block followed by
|
|
150
|
+
one block per CTA task.
|
|
151
|
+
|
|
152
|
+
The frozen `bimodal` benchmark on NVIDIA RTX A6000 `sm_86` measured medians of
|
|
153
|
+
`0.067584 ms` for pure warp, `0.070656 ms` for pure CTA, and `0.063488 ms` for
|
|
154
|
+
mixed execution. Its mixed-to-best-pure ratio is `0.939394`, passing the
|
|
155
|
+
predeclared maximum of `1.05`. This qualifies only the private canonical
|
|
156
|
+
identity-sum path. It does not add public segment syntax or public segmented
|
|
157
|
+
execution.
|
|
158
|
+
|
|
159
|
+
The research question: *can one segment-local program automatically produce
|
|
160
|
+
competitive warp, CTA, split-CTA, and persistent schedules as the runtime
|
|
161
|
+
segment-length distribution changes?* Swage did not invent ragged tensors,
|
|
162
|
+
segmented reductions, persistent kernels, or tile programming; the intended
|
|
163
|
+
contribution is the automatic derivation of the schedule from one
|
|
164
|
+
segment-local kernel.
|
|
165
|
+
|
|
166
|
+
## Getting started
|
|
167
|
+
|
|
168
|
+
```bash
|
|
169
|
+
# CPU-only: build the pinned LLVM/MLIR (once, ~1 hour), then Swage
|
|
170
|
+
./scripts/fetch_llvm.sh
|
|
171
|
+
./scripts/build_llvm.sh
|
|
172
|
+
./scripts/build_swage.sh # builds swage-opt and runs the lit suite
|
|
173
|
+
|
|
174
|
+
# Python package and tests
|
|
175
|
+
make setup
|
|
176
|
+
make test
|
|
177
|
+
```
|
|
178
|
+
|
|
179
|
+
The native bindings require the pinned LLVM/MLIR install with its Python
|
|
180
|
+
bindings enabled. Run `ninja -C build check-swage-python` after the native
|
|
181
|
+
build; it sets the build-tree `PYTHONPATH` for `mlir_swage`. On Linux with
|
|
182
|
+
PyTorch CUDA and `libcuda`, the M3 subset can then compile and launch the
|
|
183
|
+
fixed vector add. The internal native suite also qualifies M4 segmented sum
|
|
184
|
+
and max plus M5 ragged softmax. The public launch path has no CUDA toolkit
|
|
185
|
+
dependency.
|
|
186
|
+
|
|
187
|
+
See the [runnable M3 walkthrough](docs/quickstart.md#execute-fixed-vector-add).
|
|
188
|
+
A GPU is not required for the CPU and compile-only development paths.
|
|
189
|
+
|
|
190
|
+
## More
|
|
191
|
+
|
|
192
|
+
- [DESIGN.md](DESIGN.md): architecture and design invariants
|
|
193
|
+
- [ROADMAP.md](ROADMAP.md): phased plan and honest phase status
|
|
194
|
+
- [CONTRIBUTING.md](CONTRIBUTING.md): contributor paths, CPU-only onboarding
|
|
195
|
+
- [docs/adr/](docs/adr/): architecture decision records
|
|
196
|
+
|
|
197
|
+
## License
|
|
198
|
+
|
|
199
|
+
MIT. See [LICENSE](LICENSE).
|
|
@@ -0,0 +1,59 @@
|
|
|
1
|
+
# pyproject.toml
|
|
2
|
+
|
|
3
|
+
[build-system]
|
|
4
|
+
requires = ["hatchling"]
|
|
5
|
+
build-backend = "hatchling.build"
|
|
6
|
+
|
|
7
|
+
[project]
|
|
8
|
+
name = "swage-compiler"
|
|
9
|
+
dynamic = ["version"]
|
|
10
|
+
description = "Python-embedded MLIR/LLVM GPU compiler for variable-sized dense segments"
|
|
11
|
+
readme = "README.md"
|
|
12
|
+
requires-python = ">=3.10"
|
|
13
|
+
license = "MIT"
|
|
14
|
+
authors = [{ name = "Abhik Sarkar", email = "abhiksark@gmail.com" }]
|
|
15
|
+
keywords = ["mlir", "llvm", "gpu", "compiler", "cuda", "ragged", "segments"]
|
|
16
|
+
classifiers = [
|
|
17
|
+
"Development Status :: 2 - Pre-Alpha",
|
|
18
|
+
"Intended Audience :: Developers",
|
|
19
|
+
"Intended Audience :: Science/Research",
|
|
20
|
+
"Programming Language :: Python :: 3",
|
|
21
|
+
"Topic :: Software Development :: Compilers",
|
|
22
|
+
]
|
|
23
|
+
|
|
24
|
+
[project.urls]
|
|
25
|
+
Repository = "https://github.com/abhiksark/swage"
|
|
26
|
+
|
|
27
|
+
[project.optional-dependencies]
|
|
28
|
+
dev = ["build", "hatchling", "lit", "pytest", "ruff"]
|
|
29
|
+
docs = ["mkdocs", "mkdocs-material"]
|
|
30
|
+
pytorch = ["torch>=2.6"]
|
|
31
|
+
|
|
32
|
+
[tool.hatch.version]
|
|
33
|
+
path = "python/swage/__init__.py"
|
|
34
|
+
|
|
35
|
+
[tool.hatch.build.targets.wheel]
|
|
36
|
+
packages = ["python/swage"]
|
|
37
|
+
|
|
38
|
+
[tool.hatch.build.targets.sdist]
|
|
39
|
+
include = [
|
|
40
|
+
"/LICENSE",
|
|
41
|
+
"/README.md",
|
|
42
|
+
"/pyproject.toml",
|
|
43
|
+
"/python/swage",
|
|
44
|
+
]
|
|
45
|
+
|
|
46
|
+
[tool.pytest.ini_options]
|
|
47
|
+
testpaths = ["tests/python"]
|
|
48
|
+
|
|
49
|
+
[tool.ruff]
|
|
50
|
+
line-length = 80
|
|
51
|
+
target-version = "py310"
|
|
52
|
+
extend-exclude = ["build", "test/lit.cfg.py"]
|
|
53
|
+
|
|
54
|
+
[tool.ruff.lint]
|
|
55
|
+
select = ["E", "F", "W", "I", "D"]
|
|
56
|
+
ignore = ["D203", "D213"]
|
|
57
|
+
|
|
58
|
+
[tool.ruff.lint.pydocstyle]
|
|
59
|
+
convention = "google"
|