bafe 0.3.1

BAFE: Basic Algebra Fusion Engine — Fast LA with Sea-of-Nodes IR

BAFE — Basic Algebra Fusion Engine

A C++ compiler backend for linear algebra, automatic differentiation, and machine learning — built on a Sea-of-Nodes Intermediate Representation with LLVM ORCjit and MLIR integration.

Raw, heavy-lifting computational power. Stripped of fluff. Focused entirely on speed through fusion.

---

Architecture

┌─────────────────────────────────────────────────────────────────────┐
│                        Python API (pybind11)                        │
│              Lazy evaluation · Zero-copy tensors · GIL release       │
└────────────────────────────┬────────────────────────────────────────┘
                             │  (deferred execution)
┌────────────────────────────▼────────────────────────────────────────┐
│                    Sea-of-Nodes IR (Graph)                          │
│         Bidirectional edges · Shape inference · Builder API         │
└────────────────────────────┬────────────────────────────────────────┘
                             │
              ┌──────────────▼──────────────┐
              │   Optimization Pipeline     │
              │  DCE → CSE → Fusion →      │
              │  MemoryPlanner → DCE        │
              └──────────────┬──────────────┘
                             │
              ┌──────────────▼──────────────┐
              │       MLIR Dialect          │
              │  bafe → linalg/arith/math   │
              │  linalg → affine/scf (tile) │
              └──────────────┬──────────────┘
                             │
           ┌─────────────────┼─────────────────┐
           │                 │                   │
┌──────────▼──────┐ ┌───────▼────────┐ ┌───────▼────────┐
│  LLVM ORCjit    │ │  GPU (CUDA)    │ │  GPU (OpenCL)  │
│  LLJIT v2       │ │  NVRTC JIT     │ │  Fallback      │
│  Vectorized IR  │ │  Tiled kernels │ │                │
└─────────────────┘ └────────────────┘ └────────────────┘

Key Features

Sea-of-Nodes IR

No basic blocks. Every computation is a Node in a DAG with bidirectional edges (inputs + users), exposing maximum parallelism and simplifying optimization passes.

Operator Fusion (The Holy Grail)

Instead of executing sin(X) + Y as two separate GPU kernels (expensive VRAM round-trips), BAFE detects the dependency chain and fuses them into a single kernel: Out[idx] = sinf(X[idx]) + Y[idx].

Reverse-Mode Automatic Differentiation

Gradient nodes are appended to the same graph as the forward pass, so they benefit from the exact same fusion and JIT optimizations.

LLVM ORCjit v2 (On-Request-Compilation JIT)

Compile BAFE graphs to native machine code at runtime. Supports lazy compilation, vectorized loops (AVX-512/NEON), and tiled matrix multiplication.

MLIR Integration

Custom bafe dialect with progressive lowering: bafe → linalg/arith → affine/scf → llvmir.

GPU Backend

Source-to-source JIT: generate CUDA C++ kernel strings from fused subgraphs, compile via NVRTC at runtime. OpenCL fallback included. Uses dlopen for runtime CUDA loading (no build-time dependency).

Zero-Copy Python Bindings

pybind11 with Python Buffer Protocol. NumPy arrays and BAFE tensors share the same memory pointer — no copies, no serialization overhead.

---

Project Structure

BAFE/
├── CMakeLists.txt
├── include/bafe/
│   ├── ir/                    # Sea-of-Nodes IR
│   │   ├── types.h            #   DataType, Shape, Device
│   │   ├── node.h             #   Node (OpType, inputs, users)
│   │   └── graph.h            #   Graph (builder API, dump, DOT)
│   ├── optimizer/             # Optimization passes
│   │   ├── pass.h             #   Pass base class
│   │   ├── dce.h              #   Dead Code Elimination
│   │   ├── cse.h              #   Common Subexpression Elimination
│   │   ├── fusion.h           #   Operator Fusion (element-wise chains)
│   │   ├── memory_planner.h   #   In-place mutation / buffer reuse
│   │   └── pipeline.h         #   Optimization pipeline runner
│   ├── diff/                  # Automatic differentiation
│   │   └── autodiff.h         #   Reverse-mode AD (VJP rules)
│   ├── backend/               # Execution backends
│   │   ├── jit_engine.h       #   LLVM ORCjit v2 engine
│   │   ├── llvm_codegen.h     #   LLVM IR code generation
│   │   ├── gpu_codegen.h      #   CUDA/OpenCL kernel generation
│   │   ├── gpu_runtime.h      #   GPU runtime (dlopen for CUDA)
│   │   └── dispatcher.h       #   CPU/GPU dispatch
│   ├── mlir/                  # MLIR dialect
│   │   ├── bafe_dialect.h     #   'bafe' dialect registration
│   │   ├── bafe_ops.h         #   BAFE MLIR operations
│   │   ├── bafe_types.h       #   BafeTensorType
│   │   └── lowering.h         #   bafe → linalg/arith lowering
│   ├── tensor/                # Runtime tensor
│   │   └── tensor.h           #   Tensor with zero-copy support
│   └── pybind/                # Python bindings
│       └── bafe_module.h      #   pybind11 module declaration
├── src/                       # Implementation files
│   ├── ir/
│   ├── optimizer/
│   ├── diff/
│   ├── backend/
│   ├── mlir/
│   ├── tensor/
│   ├── pybind/
│   └── demo.cpp               # End-to-end demo
├── python/bafe/
│   ├── __init__.py            # Python API (lazy Tensor, Graph, nn)
│   └── nn.py                  # Neural network building blocks
└── tests/

---

Building

Prerequisites

• C++17 compiler (GCC 9+, Clang 10+)
• CMake 3.14+
• (Optional) LLVM + MLIR development headers
• (Optional) CUDA Toolkit (for GPU backend)
• (Optional) pybind11 (for Python bindings)

Quick Build (Core only — no LLVM/MLIR required)

mkdir build && cd build
cmake .. -DCMAKE_BUILD_TYPE=Release
make -j$(nproc)
./bafe_demo

Full Build (with LLVM + MLIR)

mkdir build && cd build
cmake .. \
  -DBAFE_ENABLE_MLIR=ON \
  -DBAFE_ENABLE_LLVM=ON \
  -DBAFE_ENABLE_PYBIND=ON \
  -DLLVM_DIR=$(llvm-config --cmakedir) \
  -DMLIR_DIR=$(llvm-config --cmakedir)/../mlir
make -j$(nproc)

---

Quick Start (C++)

#include "bafe/ir/graph.h"
#include "bafe/optimizer/pipeline.h"
#include "bafe/diff/autodiff.h"

bafe::Graph graph;

// Build: loss = reduce_sum(relu(matmul(W, x) + b))
auto* W = graph.new_parameter(bafe::DataType::Float32, {4, 3}, "W");
auto* x = graph.new_parameter(bafe::DataType::Float32, {3, 1}, "x");
auto* b = graph.new_parameter(bafe::DataType::Float32, {4, 1}, "b");

auto* Wx   = graph.matmul(W, x);
auto* Wx_b = graph.add(Wx, b);
auto* act  = graph.relu(Wx_b);
auto* loss = graph.reduce_sum(act);
graph.mark_output(loss);

// Optimize (DCE → CSE → Fusion → MemoryPlan → DCE)
auto pipeline = bafe::Pipeline::default_pipeline();
pipeline.run(graph);

// Compute gradients (reverse-mode AD)
bafe::reverse_mode_autodiff(graph, loss);
auto* dW = bafe::grad(graph, W, loss);

// Visualize
graph.dump_dot_to_file("graph.dot");

Quick Start (Python)

import bafe
import numpy as np

x = bafe.Tensor(np.array([1.0, 2.0, 3.0]))
y = bafe.Tensor(np.array([4.0, 5.0, 6.0]))

z = (x + y).relu()          # lazy: recorded in graph
result = z.numpy()           # triggers compilation + execution

# Neural networks
from bafe.nn import Linear, ReLU, Sequential

model = Sequential(Linear(784, 256), ReLU(), Linear(256, 10))
output = model(x)

---

Optimization Pipeline

Pass	Purpose
DCE	Walk backward from outputs, delete unreachable nodes
CSE	Hash nodes by (OpType, inputs, dtype, shape), replace duplicates
Fusion	Group consecutive element-wise ops into FusedGroup meta-nodes
MemoryPlanner	Analyze node lifetimes, reuse buffers for dead tensors

The fusion pass is the performance killer feature. It detects chains like:

X → Sin → Add(Y) → ReLU → Output

and merges them into a single FusedGroup that generates one GPU kernel:

__global__ void fused_kernel(float* X, float* Y, float* Out, int N) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < N) {
        float v = sinf(X[idx]) + Y[idx];
        Out[idx] = v > 0.0f ? v : 0.0f;
    }
}

---

Automatic Differentiation

BAFE implements reverse-mode AD (backpropagation) using VJP (Vector-Jacobian Product) rules:

Forward Op	Gradient Rule
Add(a, b)	grad_a += grad_out, grad_b += grad_out
Mul(a, b)	grad_a += grad_out * b, grad_b += grad_out * a
MatMul(a, b)	grad_a += grad_out @ b^T, grad_b += a^T @ grad_out
ReLU(a)	grad_a += grad_out * (a > 0)
Exp(a)	grad_a += grad_out * exp(a)
Sin(a)	grad_a += grad_out * cos(a)

Gradient nodes are appended to the same graph, so they get fused and JIT-compiled too.

---

MLIR Dialect

The bafe MLIR dialect provides progressive lowering:

bafe.add     →  arith.addf
bafe.mul     →  arith.mulf
bafe.matmul  →  linalg.matmul
bafe.relu    →  arith.maximum + arith.constant(0)
bafe.fused_group  →  inlined element-wise chain

Then further lowered:

linalg.matmul  →  affine.for (tiled)  →  scf.for  →  llvmir

---

Tech Stack

Component	Purpose
LLVM ORCjit v2	JIT compilation to native CPU code (vectorized, tiled)
MLIR	Multi-level IR transformation (bafe → linalg → affine → llvm)
NVRTC	Runtime CUDA compilation (source-to-source JIT)
pybind11	Zero-copy Python bindings with buffer protocol

---

License

See LICENSE.