Skip to content

ky2renzz/bend-ml

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

3 Commits
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

Bend-ML

GitHub stars Release License Python PyPI

PyTorch-to-Bend converter for massive GPU parallelism via HVM2 Interaction Nets

Convert PyTorch models to Bend code that runs on all CPU cores / GPUs with automatic parallelization

🎬 Demo

# PyTorch model
model = nn.Sequential(
    nn.Linear(784, 256),
    nn.ReLU(),
    nn.Linear(256, 10)
)

# One line to Bend
code = pytorch_to_bend(model)

# Run on GPU (massively parallel!)
# bend run-cu model.bend

Generated Bend code uses tree-fold parallelism — O(log n) depth on HVM2:

def matrix_tree_mul(vec, matrix):
  match matrix:
    case MCons(row1, MCons(row2, MNil)):
      return Cons(vec_dot(vec, row1), Cons(vec_dot(vec, row2), Nil))
    case _:
      (left, right) = matrix_split(matrix)
      vec_append(matrix_tree_mul(vec, left), matrix_tree_mul(vec, right))

How It's Different

Aspect Naive Converters bend-ml
Data structures Raw lists [f24] data Vector, data Matrix
Operations List-based (dot_product on [f24]) ADT-based (vec_dot on Vector)
Matrix-Vector Sequential tail-rec (matvec_fold) Tree-fold parallel (matrix_tree_mul)
Reduction depth O(n) sequential O(log m × log n) tree-fold
Fan-out Sequential row processing Parallel matrix split + divide-and-conquer
Control flow if/else chains switch + pattern matching on ADTs
Model size Truncated/fake limits Full models with all weights
Style Haskell-like lists Bend with Interaction Net ADTs

Installation

pip install bend-ml

Or from source:

git clone https://github.com/ky2renzz/bend-ml
cd bend-ml
pip install -e .

Requirements

  • Python 3.8+
  • PyTorch 2.0+
  • Bend (optional, for running generated code): cargo install bend-lang

Quick Start

import torch
import torch.nn as nn
from bend_ml import pytorch_to_bend

class SimpleNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(10, 32)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(32, 2)

    def forward(self, x):
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        return x

model = SimpleNet()
code = pytorch_to_bend(model)
print(code)

Generated Bend code:

# Algebraic data types
data Vector:
  | Nil
  | Cons { head: f24, tail: Vector }

data Matrix:
  | MNil
  | MCons { row: Vector, rows: Matrix }

# Tree-fold sum - O(log n) depth
def vec_tree_sum(v: Vector) -> f24:
  match v:
    case Nil: return 0.0
    case Cons(x, Nil): return x
    case Cons(x, Cons(y, Nil)): return x + y
    case _:
      (left, right) = vec_split(v)
      lsum = vec_tree_sum(left)
      rsum = vec_tree_sum(right)
      return lsum + rsum

# Parallel dot product via tree-fold
def vec_dot(a: Vector, b: Vector) -> f24:
  products = vec_zip_mul(a, b)
  vec_tree_sum(products)

# Matrix-Vector multiplication - O(log m × log n) depth
def matrix_tree_mul(vec: Vector, matrix: Matrix) -> Vector:
  match matrix:
    case MNil: return Nil
    case MCons(row, MNil):
      dot = vec_dot(vec, row)
      return Cons(dot, Nil)
    case MCons(row1, MCons(row2, MNil)):
      dot1 = vec_dot(vec, row1)
      dot2 = vec_dot(vec, row2)
      return Cons(dot1, Cons(dot2, Nil))
    case _:
      (left, right) = matrix_split(matrix)
      left_result = matrix_tree_mul(vec, left)
      right_result = matrix_tree_mul(vec, right)
      vec_append(left_result, right_result)

# Linear layer
def linear_fc1(input: Vector) -> Vector:
  weight_matrix = list_to_matrix(WEIGHT_fc1)
  outputs = matrix_tree_mul(input, weight_matrix)
  vec_zip_add(outputs, list_to_vec(BIAS_fc1))

# Forward pass with I/O conversion + pure Vector composition
def forward(input: [f24]) -> [f24]:
  input_vec = list_to_vec(input)
  output_vec = linear_fc2(tanh_tanh(linear_fc1(input_vec)))
  vec_to_list(output_vec)

Supported Layers

Layer Status Pattern Parallelism
nn.Linear Full matrix_tree_mul O(log M × log N) tree-fold
nn.ReLU Full vec_tree_map with relu O(log N) tree-fold
nn.LeakyReLU Full vec_tree_map with leaky_relu O(log N) tree-fold
nn.GELU Full vec_tree_map with gelu O(log N) tree-fold
nn.SiLU Full vec_tree_map with silu O(log N) tree-fold
nn.Sigmoid Full vec_tree_map with sigmoid O(log N) tree-fold
nn.Tanh Full vec_tree_map with tanh O(log N) tree-fold
nn.Softmax Full vec_tree_fold_max + vec_tree_fold_sum O(log N) tree-fold
nn.Conv2d Full im2col + tensor4d_tree_fold output channels O(log C × log N) tree-fold
nn.BatchNorm2d Full tensor4d_tree_zip_map channels O(log C) tree-fold

Examples

# Run tests
python tests/test_converter.py

# MLP for MNIST - generates full model
python examples/mlp_mnist.py

# Simple classifier - small model for testing
python examples/simple_classifier.py

# Modern CNN with GELU, SiLU, BatchNorm
python examples/resnet_example.py

How It Works

1. Algebraic Data Types

def process(xs: [f24]) -> [f24]

data Vector:
  | Nil
  | Cons { head: f24, tail: Vector }

def process(v: Vector) -> Vector
  match v:
    case Nil: return Nil
    case Cons(h, t): return Cons(f(h), process(t))

2. Tree-Fold Parallelism

def bad_sum(xs: [f24]) -> f24:
  match xs:
    case []: return 0
    case [x, *rest]: return x + bad_sum(rest)

def vec_tree_sum(v: Vector) -> f24:
  match v:
    case Cons(x, Cons(y, Nil)): return x + y
    case _:
      (left, right) = vec_split(v)
      lsum = vec_tree_sum(left)
      rsum = vec_tree_sum(right)
      return lsum + rsum

3. Matrix-Vector with Tree-Fold

def matvec_fold(vec: [f24], matrix: [[f24]], acc: [f24]) -> [f24]:
  match matrix:
    case []: return acc
    case [row, *rows]:
      dot = dot_product(vec, row)
      matvec_fold(vec, rows, [*acc, dot])

def matrix_tree_mul(vec: Vector, matrix: Matrix) -> Vector:
  match matrix:
    case MNil: return Nil
    case MCons(row1, MCons(row2, MNil)):
      dot1 = vec_dot(vec, row1)
      dot2 = vec_dot(vec, row2)
      return Cons(dot1, Cons(dot2, Nil))
    case _:
      (left, right) = matrix_split(matrix)
      left_result = matrix_tree_mul(vec, left)
      right_result = matrix_tree_mul(vec, right)
      vec_append(left_result, right_result)

4. Pure Function Composition

def forward(input: [f24]) -> [f24]:
  linear_fc2(tanh_tanh(linear_fc1(input)))

5. Tree-fold for All Data Structures

# Tree-fold on Lists (parallel element-wise operations)
def list_tree_map(xs: [f24], f: f24 -> f24) -> [f24]:
  match xs:
    case []: return []
    case [x]: return [f(x)]
    case [x, y]: return [f(x), f(y)]
    case _:
      (left, right) = list_split(xs)
      lmap = list_tree_map(left, f)
      rmap = list_tree_map(right, f)
      list_append(lmap, rmap)

# Tree-fold on Tensor4D (parallel channel processing)
def tensor4d_tree_map(t: Tensor4D, f: [[f24]] -> [[f24]]) -> Tensor4D:
  match t:
    case TNil: return TNil
    case TCons(ch, TNil): return TCons(f(ch), TNil)
    case TCons(ch1, TCons(ch2, TNil)):
      return TCons(f(ch1), TCons(f(ch2), TNil))
    case _:
      (left, right) = tensor4d_split(t)
      left_map = tensor4d_tree_map(left, f)
      right_map = tensor4d_tree_map(right, f)
      tensor4d_append(left_map, right_map)

6. Conv2d with Tree-fold Channel Processing

# Tree-fold generation of output channels O(log C) depth
def conv2d_layer_output_tree(input, cols, start_idx, end_idx, ...) -> Tensor4D:
  switch end_idx - start_idx:
    case 0: return TNil
    case 1:
      return TCons(compute_single(...), TNil)
    case 2:
      ch1 = compute_single(input, cols, start_idx, ...)
      ch2 = compute_single(input, cols, start_idx + 1, ...)
      return TCons(ch1, TCons(ch2, TNil))
    case _:
      mid = (start_idx + end_idx) / 2
      left = conv2d_layer_output_tree(input, cols, start_idx, mid, ...)
      right = conv2d_layer_output_tree(input, cols, mid, end_idx, ...)
      tensor4d_append(left, right)

Running Generated Code

# Install Bend
cargo install bend-lang

# Sequential (for testing)
bend run-rs model.bend

# Parallel (C interpreter, uses all cores)
bend run-c model.bend

# Massively parallel (CUDA, requires NVIDIA GPU)
bend run-cu model.bend

Performance Characteristics

Metric Naive List-based bend-ml
Data structures Raw lists [f24] ADTs: Vector, Matrix, Tensor4D
Vector reduction O(n) sequential O(log n) tree-fold on ADT
List operations O(n) sequential O(log n) tree-fold
Dot product O(n) on lists O(log n) on Vector ADT
Matrix-vector O(n×m) sequential fold O(log m × log n) tree-fold
Conv2d output channels O(C) sequential O(log C) tree-fold
BatchNorm2d channels O(C) sequential O(log C) tree-fold
Parallelism Sequential accumulation Divide-and-conquer on ADTs
Interaction Nets List overhead Optimal ADT reduction
Scalability 1 core All cores/GPUs via HVM2

HVM2's Interaction Nets provide automatic parallelism. Tree-fold patterns create independent redexes that HVM2 reduces in parallel.

Project Structure

bend-ml/
├── bend_ml/
│   ├── __init__.py
│   └── converter.py      │                           │                           ├── tests/
│   └── test_converter.py├── examples/
│   ├── mlp_mnist.py      │   └── simple_classifier.py  # Compact example with data types
└── README.md

License

MIT

Acknowledgments

  • Bend by HigherOrderCO
  • HVM2 — Higher-order Virtual Machine 2 with Interaction Nets
  • Victor Taelin for Interaction Combinators

About

PyTorch to Bend converter - transforms neural networks into parallel functional code using Interaction Nets for HVM2

Topics

Resources

License

Stars

1 star

Watchers

0 watching

Forks

Packages

 
 
 

Contributors

Languages