Skip to content
AriBandyoPublic

About

No description, website, or topics provided.

Resources

Stars

0 stars

Watchers

0 watching

Forks

Repository files navigation

Coalesce

A fusion-driven GPU kernel compiler.

Coalesce lowers a high-level fused-operator dialect through MLIR to hand-tuned GPU kernels, beating the equivalent eager PyTorch op sequence.

The name carries both meanings deliberately: coalesced memory access (the core GPU performance term) and coalescing operations together (operator fusion). A library ships fixed, unfused kernels — it cannot fuse your GEMM + bias + GELU into one kernel and eliminate the memory round trips. That fusion is the reason this project exists.

This is a benchmark-driven project. Every phase deliverable is a measured comparison against a baseline. The numbers are the point: how much faster, and why.


Project structure

The project has two halves, in a locked order. Kernel first, then compiler.

Half Phases What it is
Kernel 0–2 Hand-write CUDA kernels and make them fast. The tuned-kernel backend.
Compiler 3–4 Make an MLIR frontend generate those kernels automatically.

Kernel-first is intentional: it produces something portfolio-visible in weeks, then deepens into the part most applicants won't have.


Status

Phase Description State
0 First measured kernels — vector add, naive matmul ✅ Complete
1 Fused GEMM + bias + GELU — the anchor ✅ Complete
1A Tiled matmul, benchmarked ✅ (numbers superseded — see Phase 2)
1B Bottleneck analysis ✅ Analytical (counters unavailable — see notes)
1C Fused kernel + scaling study ✅ 4.01% at N=4096, bit-exact
1D Writeup ✅ docs/phase1.md
2 Register tiling + corrected harness + shape sweep ✅ docs/phase2.md
2A Benchmark harness rebuild (L2 flush, median timing) ✅ bench/bench_utils.cuh
2B 4×4 register-tiled kernel ✅ 3849 GFLOPS @ 4096²
2C cuBLAS baseline + 512–4096 shape sweep ✅ bench/results/phase2.csv
2.5 Triton port — compiler vs. hand-tuned, tile sweep ✅ 05_matmul_triton/
3 MLIR dialect + lowering passes 🔨 In progress
3A Dialect scaffolding + coalesce-opt driver + lit tests ✅
3B Fusion rewrite pass (matmul → relu, epilogue attribute) ✅ 3 lit tests
3C Library-call lowering to extern "C" entry points ✅ 2 lit tests
3D End-to-end demo, fused vs. unfused ◻ Next
3E lit/FileCheck tests + docs/phase3.md 🔨 tests ✅, writeup pending
4 Autotuning + capstone writeup ◻ Not started

Results

All CUDA measurements FP32, square matrices, median-of-100, L2 flushed between timed runs, after boost-clock stabilisation.

Kernel progression

N naive tiled regtile cuBLAS regtile vs tiled % of cuBLAS
512 759.8 989.2 2880.7 3692.2 2.91× 78%
1024 781.6 1021.5 4378.2 5115.0 4.29× 86%
2048 791.1 1034.7 4504.0 5981.2 4.35× 75%
4096 611.1 822.3 3848.9 5576.4 4.68× 69%

GFLOPS. Reproduce with cd bench && ./bench_all; raw output in bench/results/phase2.csv.

The N=4096 column is the defensible one. At smaller sizes the operand working set fits in the 32 MB L2, so those figures are cache-assisted — see docs/phase2.md §2.

max_diff against the tiled kernel is 0.000e+00 at every size for both naive and regtile.

The Phase 1 numbers were corrected

Phase 1 reported the tiled kernel at 1031 GFLOPS against an analytical DRAM ceiling it should not have been able to exceed. The harness was at fault: no L2 flush between iterations, so runs 2–100 read operands from a warm cache, and a single event pair around the whole loop made a mean the only available statistic.

kernel Phase 1 (N=1024) corrected change
naive 646 781.6 +21%
tiled 1031 1021.5 −1%

The naive figure moved further than the tiled one, and in the opposite direction — the old 5-iteration warm-up had been timing a GPU below its boost clock. Full account in docs/phase2.md §1.

Register tiling

Each thread owns a 4×4 patch of C in registers, turning the inner loop into a rank-1 outer product: 8 shared-memory reads per 16 FMAs instead of 2 reads per 1 FMA — a 4× gain in arithmetic intensity for identical shared-memory traffic.

BM = 64, BN = 64, BK = 8, TM = TN = 4
threads = 256, shared = 4 KB
ptxas: 62 registers, 0 spills → 4 blocks/SM, 67% occupancy

Registers, not shared memory, are the binding constraint on this configuration.

Fusion study (03_matmul_fused)

GEMM + bias + GELU, fused into one kernel vs. the same matmul followed by a separate elementwise kernel.

N Intermediate size vs. 32 MB L2 Fused Unfused Speedup
1024 4 MB fits easily 2.0881 ms 2.1079 ms 1.009× (0.95%)
4096 64 MB 2× over 146.13 ms 151.99 ms 1.040× (4.01%)

Both sizes: max |fused − unfused| = 0.000e+00 — bit-identical, not merely close.

The finding: fusion's value depends on whether the intermediate escapes L2. When it fits in cache, the round trip you eliminate was already cheap.

⚠️ These figures were measured under the pre-Phase-2 harness (mean timing, no L2 flush) and have not been re-run. Both sides of the comparison were measured identically, so the ratio is expected to hold; the absolute millisecond figures are not comparable to the table above. Re-measurement through bench_all is pending.

Triton port (05_matmul_triton)

The same tiled matmul written in Triton instead of CUDA, to measure what a compiler produces against a hand-tuned kernel on identical hardware.

Harness flushes L2 (256 MB buffer) before every timed iteration and reports the median, matching bench_utils.cuh. Correctness uses relative Frobenius norm rather than elementwise allclose: random products contain near-zero entries where the relative term vanishes and an absolute tolerance decides, failing the whole check on one element even when the kernel is correct.

Kernel Precision GFLOPS @ 4096²
Hand-written CUDA, register-tiled FP32 3849
Triton, 64×64×16 FP32 ~4200
Triton, 64×64×16 tf32 6411
Triton, 128×128×32, 8 warps tf32 8794
torch / cuBLAS tf32 ~9460

Relative error: 4.2e-07 at FP32, 7.7e-04 at tf32 — three orders of magnitude apart, consistent with tf32's 10-bit mantissa.

Occupancy is not throughput. Triton's kernel used 96 registers per thread against 62 for the hand-written version, both with zero spills — roughly 5 resident blocks per SM against 8. The lower-occupancy kernel matched or beat the higher-occupancy one, trading occupancy for more loads in flight. It also allocated 16 KB of shared memory (two pipeline stages) rather than the three requested.

Tile shape was worth more than anything else measured. The 64×64×16 shape derived by hand in Phase 2 was the slowest working config in the sweep. Moving to 128×128×32 gained 37% on identical kernel source. 128×128×64 failed to launch at all, requiring 196,608 bytes of shared memory against a 101,376-byte limit — the constraint that bounds the search.

Run-to-run variance on a laptop 4060 is roughly 10% from thermal throttling, so the CUDA-vs-Triton FP32 gap sits inside the noise band and is reported as comparable rather than a win either way. The tile-shape and tensor-core gaps are far outside it.

MLIR pipeline (Phase 3)

coalesce-opt, a compiler driver built on MLIR, with the coalesce dialect registered alongside MLIR's built-in dialects:

$ ./build/bin/coalesce-opt --help
OVERVIEW: Coalesce optimizer driver
Available Dialects: arith, builtin, coalesce, func, memref

Two passes, composable in one invocation:

./build/bin/coalesce-opt input.mlir \
  --coalesce-fuse-epilogue \
  --coalesce-lower-to-library

A separate coalesce.matmul and coalesce.relu become a single call into an extern "C" entry point, with the intermediate buffer eliminated:

module {
  func.func private @coalesce_matmul_relu(memref<4x4xf32>, memref<4x4xf32>, memref<4x4xf32>)
  func.func @f(%arg0: memref<4x4xf32>, %arg1: memref<4x4xf32>, %arg2: memref<4x4xf32>) {
    call @coalesce_matmul_relu(%arg0, %arg1, %arg2) : (memref<4x4xf32>, memref<4x4xf32>, memref<4x4xf32>) -> ()
    return
  }
}

Fusion legality is guarded on four conditions — unique producer, same block, producer precedes consumer, and no other user of the intermediate buffer — with negative lit tests for each failure mode. Design rationale in mlir/README.md; the build log with every error and fix is in docs/phase3.md.

Derived from mlir/examples/standalone in the LLVM tree (Apache 2.0 with LLVM exceptions). Test suite: 6 passing, 1 unsupported (Python bindings, disabled).

⚠️ coalesce_matmul_relu is declared by the compiler and implemented by nothing. Linking the pipeline output fails until 3D writes the wrapper around the Phase 2 kernel. The compiler infrastructure is the artifact so far; the end-to-end number is what 3D earns.


Repository layout

01_matmul/          naive.cuh    + matmul.cu           naive
02_matmul_tiled/    tiled.cuh    + matmul_tiled.cu     shared-memory tiled
03_matmul_fused/    fused.cuh    + matmul_fused.cu     GEMM + bias + GELU
04_matmul_regtile/  regtile.cuh  + matmul_regtile.cu   4×4 register tiled
05_matmul_triton/   matmul_triton.py                   Triton port + tile sweep
bench/              bench_utils.cuh, bench_all.cu      harness + driver
mlir/               include/Coalesce, lib, coalesce-opt   MLIR dialect + passes
docs/               phase1.md, phase2.md, phase3.md

Kernels live in .cuh headers alongside a host launcher; each .cu holds only a main(). This is what lets one driver benchmark every kernel in a single process against one set of buffers — previously each kernel was its own program with its own timing code, and cross-kernel comparisons assumed those harnesses were equivalent.


Hardware / environment

GPU NVIDIA GeForce RTX 4060 Laptop, 8 GB
Architecture Ada Lovelace, sm_89
FP32 peak ~14.6 TFLOPS
DRAM bandwidth ~272 GB/s (128-bit GDDR6 @ 17 Gbps)
L2 cache 32 MB
Registers per SM 65,536
Max threads per SM 1,536 (48 warps)
Driver 610.88
CUDA toolkit 12.8
Host WSL2, Ubuntu 25.10
Host compiler g++-14 (via -ccbin)
Triton 3.8.0, source build at ~/triton/triton
PyTorch 2.13.0+cu130
LLVM / MLIR 23.0.0git, vendored via Triton

Build and run

CUDA kernels

The full benchmark sweep — every kernel, every shape, one process:

cd bench
nvcc -arch=sm_89 -ccbin g++-14 -O3 -Xptxas -v bench_all.cu -o bench_all -lcublas
./bench_all

Individual kernels, for debugging in isolation:

cd 02_matmul_tiled   && nvcc -ccbin g++-14 -arch=sm_89 -O3 matmul_tiled.cu   -o matmul_tiled   && ./matmul_tiled
cd 03_matmul_fused   && nvcc -ccbin g++-14 -arch=sm_89 -O3 matmul_fused.cu   -o matmul_fused   && ./matmul_fused
cd 04_matmul_regtile && nvcc -ccbin g++-14 -arch=sm_89 -O3 matmul_regtile.cu -o regtile        && ./regtile

-ccbin g++-14 is required — the default host compiler on Ubuntu 25.10 is newer than CUDA 12.8 supports. -lcublas only for bench_all.

-Xptxas -v reports register and shared-memory usage per kernel. Worth leaving on: a register spill in the register-tiled kernel would silently make it slower than the tiled one.

Triton port

cd 05_matmul_triton
source ~/triton/triton/.venv/bin/activate
python matmul_triton.py                                  # FP32, correctness + bench + PTX
python matmul_triton.py --precision tf32 --no-ptx        # tensor cores
python matmul_triton.py --precision tf32 --sweep --no-ptx  # tile-shape sweep

MLIR dialect

Builds against the LLVM 23.0.0git that ships with Triton, so no separate LLVM build is required.

cd mlir
LLVM_DIR=$(ls -d ~/.triton/llvm/llvm-*-ubuntu-x64*)
cmake -G Ninja -B build \
  -DMLIR_DIR="$LLVM_DIR/lib/cmake/mlir" \
  -DMLIR_INCLUDE_TESTS=ON \
  -DLLVM_EXTERNAL_LIT="$(which lit)" \
  -DCMAKE_BUILD_TYPE=Release \
  -DCMAKE_C_COMPILER="$LLVM_DIR/bin/clang" \
  -DCMAKE_CXX_COMPILER="$LLVM_DIR/bin/clang++"
cmake --build build --target check-coalesce

Use the clang that ships with the vendored LLVM rather than the system g++: the MLIR libraries were built with it, and mixing toolchains across a C++ library boundary risks ABI mismatches that surface as undefined-symbol walls at link time.

Triton's vendored LLVM does not ship the count test utility that lit expects. A replacement lives in mlir/tools-shim/count and must be copied into $LLVM_DIR/bin/ before the test target will run.


Notes for future sessions

  • Benchmarks require a timed warm-up. GPU Boost causes ~30% run-to-run variance otherwise. Every benchmark runs the kernel continuously until 1000 ms of GPU work has elapsed before opening the timed window. Do not remove this.
  • Benchmarks require an L2 flush between iterations. A 64 MB buffer (2× L2) is cudaMemsetAsync'd before each timed run. Without it, a kernel whose working set fits in the 32 MB L2 reports above its own DRAM roofline. This is what invalidated the Phase 1 numbers.
  • Time each iteration separately, report the median. One event pair around the whole loop makes a mean the only statistic available, and one throttled iteration then moves the reported figure.
  • Thermal throttling sets a ~10% noise floor. Repeated runs of the identical Triton config at N=4096 spread across 3910–4282 GFLOPS. Do not report a difference smaller than that as a result.
  • Test data must be able to expose a broken kernel. A=1, B=1 hides transposed indices — the wrong element holds the same value as the right one. The harness uses ((i*7)%13)*0.1f and ((i*3)%11)*0.1f, checked against the tiled kernel's output. For the fused kernel, B = 1/N so the accumulator is 1.0 and GELU sees values where the curve actually bends.
  • Elementwise allclose is the wrong correctness test for reduced precision. Random products contain near-zero entries where the relative term vanishes and atol alone decides, failing on one element while the kernel is fine. Use a relative norm ratio.
  • A kernel-name collision produces wrong numbers, not a build error. With all kernel headers included in one translation unit, a launcher calling matmul instead of matmul_tiled compiles cleanly and silently benchmarks the wrong kernel. Caught only because max_diff was exactly zero where it had no business being.
  • Nsight Compute hardware counters do not work on this machine. Driver 610.88 uses capability-based permissions; WSL2 routes GPU access through /dev/dxg and has no /dev/nvidia-caps to grant. Not fixable from the Linux side, and sudo does not help. Bottleneck analysis is done analytically instead — the method is documented in docs/phase1.md and is arguably more instructive.
  • Use erff, not the tanh approximation, for GELU. PyTorch's default is the exact erf form. Mixing them produces small mismatches that look like kernel bugs.
  • lit declares tool dependencies in two places. test/CMakeLists.txt controls build ordering; lit's own add_tool_substitutions decides what must exist on disk. Both have to agree, and a tool missing from the second is a fatal error regardless of whether any test uses it.
  • In MLIR, read only the first error from generated code. TableGen emits one enormous class declaration per op, so a single unresolved name produces a page of cascading noise. Scrolling to the bottom is the wrong instinct.
  • [N/N] Linking with no compile steps means a file is missing from CMake. An unchanged link error after adding a .cpp almost always means the source list was never updated.

About

No description, website, or topics provided.

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages